"""
jis-core OSAPI — bootstrap handshake (v1.0 protocol, §3).

Reference: docs/specs/osapi-protocol-v1.md
Intentional mirror of tibet-core/osapi.py — same wire-protocol, different
default endpoints and (v0.5.0b1+) different ops (claim / bind / fira instead of
emit / query / fork). The handshake itself is generic by design so packages can
bootstrap to both OSAPI's with one code-path.

Scope of this module (v0.5.0a1):
  - HELLO / WELCOME / ACK handshake over line-delimited JSON
  - UDS + TCP transport
  - JIS_OSAPI_URL discovery + TIBET_SOFT_BOOTSTRAP fallback (shared env-var)
  - Minimal jis-OSAPI server (reference implementation for tests + dev)
  - Error types per spec (BootstrapError, IdentityError, VersionError, ProtocolError)

NOT in this module yet (v0.5.x roadmap):
  - claim/bind/fira operations (§4) — Session is a handle only for now
  - heartbeat loop (§5)
  - soft-stop subscriber (§6)
  - Real Ed25519 actor_claim verification — uses the compiled Rust kernel
    (jis_core.jis_core) once v0.4.0 lands the verify hook
"""
from __future__ import annotations

import hashlib
import hmac
import json
import os
import secrets
import socket
import sys
import threading
import time
from dataclasses import dataclass, field
from typing import Callable, Optional

# ── Protocol constants (jis-OSAPI flavor) ───────────────────────────────────

PROTOCOL_VERSION = "1.0"
SUPPORTED_VERSIONS = ("1.0",)
DEFAULT_HEARTBEAT_INTERVAL_MS = 30_000
SOFT_STOP_TOPIC = "tibet.cap-bus.events"

DEFAULT_UDS_PATH = "/var/run/jis/osapi.sock"
DEFAULT_TCP_HOST = "127.0.0.1"
DEFAULT_TCP_PORT = 18444   # tibet-OSAPI = 18443, jis-OSAPI = 18444


# ── Errors (per spec §3) ────────────────────────────────────────────────────

class OSAPIError(Exception):
    """Base for all OSAPI protocol errors."""


class BootstrapError(OSAPIError):
    """OSAPI unreachable, or handshake failed and no soft-bootstrap allowed."""


class IdentityError(OSAPIError):
    """actor_claim invalid (missing, malformed, or signature-failed)."""


class VersionError(OSAPIError):
    """Protocol-version mismatch; no compatible version negotiated."""


class ProtocolError(OSAPIError):
    """Malformed line, unexpected message kind, or sequence-violation."""


# ── URL parsing (mirror of tibet-core) ──────────────────────────────────────

def _resolve_url(explicit: Optional[str]) -> str:
    """Discovery order (spec §2): explicit arg > JIS_OSAPI_URL > well-known UDS > TCP fallback."""
    if explicit:
        return explicit
    env = os.environ.get("JIS_OSAPI_URL")
    if env:
        return env
    if os.path.exists(DEFAULT_UDS_PATH):
        return f"unix://{DEFAULT_UDS_PATH}"
    return f"tcp://{DEFAULT_TCP_HOST}:{DEFAULT_TCP_PORT}"


def _connect(url: str, timeout: float = 5.0) -> socket.socket:
    if url.startswith("unix://"):
        path = url[len("unix://"):]
        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        s.settimeout(timeout)
        s.connect(path)
        return s
    if url.startswith("tcp://"):
        url = url[len("tcp://"):]
    host, _, port = url.partition(":")
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.settimeout(timeout)
    s.connect((host, int(port or DEFAULT_TCP_PORT)))
    return s


# ── Wire protocol (line-delimited JSON) ─────────────────────────────────────

def _send(sock: socket.socket, msg: dict) -> None:
    sock.sendall((json.dumps(msg) + "\n").encode("utf-8"))


def _recv_line(sock: socket.socket, max_bytes: int = 64 * 1024) -> dict:
    buf = bytearray()
    while True:
        chunk = sock.recv(4096)
        if not chunk:
            raise ProtocolError("connection closed before line terminator")
        buf.extend(chunk)
        if b"\n" in buf:
            line, _ = bytes(buf).split(b"\n", 1)
            break
        if len(buf) > max_bytes:
            raise ProtocolError(f"line exceeds {max_bytes} bytes without terminator")
    try:
        return json.loads(line)
    except json.JSONDecodeError as e:
        raise ProtocolError(f"malformed JSON line: {e}") from None


def _ack_hmac(session_token: str) -> str:
    return hmac.new(session_token.encode("utf-8"), b"ack-echo", hashlib.sha256).hexdigest()


def _new_session_token() -> str:
    """jis-OSAPI uses `jis_` prefix to distinguish from tibet-OSAPI session-tokens."""
    return f"jis_{int(time.time() * 1_000_000):d}_{secrets.token_hex(4)}"


# ── Session handle (client-side, post-handshake) ────────────────────────────

@dataclass
class Session:
    actor: str
    session_token: str
    chain_handle: str          # for cross-OSAPI correlation (will mirror tibet-side)
    identity_binding: str      # primary handle for jis-OSAPI ops
    heartbeat_interval_ms: int = DEFAULT_HEARTBEAT_INTERVAL_MS
    soft_stop_topic: str = SOFT_STOP_TOPIC
    server_url: str = ""
    is_soft_bootstrap: bool = False
    _local_state: dict = field(default_factory=dict)   # populated only in soft-bootstrap mode

    # ── Operations (spec §4 jis-OSAPI) ──────────────────────────────────────

    def claim(self, claim_data: dict) -> dict:
        """Submit a JIS claim for verification. Returns {verdict, denials}.

        claim_data shape: {actor, clearance, role?, department?, geo?, ...}
        """
        if self.is_soft_bootstrap:
            return _local_claim(self, claim_data)
        return _call_op(self, "claim", claim=claim_data)

    def bind(self, *, actor: str, intent: str, ttl_seconds: int = 3600) -> dict:
        """Bind an identity to an action/intent. Returns {binding_token, expires_at}."""
        if self.is_soft_bootstrap:
            return _local_bind(self, actor, intent, ttl_seconds)
        return _call_op(self, "bind", actor=actor, intent=intent, ttl_seconds=ttl_seconds)

    def fira(self, *, peer: str, phase: str) -> dict:
        """First-Initiation Revoke/Accept peer-trust handshake. Returns {trust, next_phase}.

        phase ∈ {"F", "R", "A"}.
        """
        if self.is_soft_bootstrap:
            return _local_fira(self, peer, phase)
        return _call_op(self, "fira", peer=peer, phase=phase)

    def close(self) -> None:
        """v0.5.x: no-op. v0.6.x: send BYE, stop heartbeat thread."""
        pass


# ── Op-call helper (request-response per op over a fresh socket) ────────────

def _call_op(session: "Session", op: str, **payload) -> dict:
    """Open socket, send OP-message, get response, close. (Mirrors tibet-core.)"""
    sock = _connect(session.server_url)
    try:
        msg = {"kind": "OP", "op": op, "session_token": session.session_token}
        msg.update({k: v for k, v in payload.items() if v is not None})
        _send(sock, msg)
        resp = _recv_line(sock)
        if resp.get("kind") == "ERROR":
            code = resp.get("code", "PROTOCOL")
            detail = resp.get("detail", "")
            raise OSAPIError(f"{code}: {detail}")
        return resp
    finally:
        try:
            sock.close()
        except Exception:
            pass


# ── Policy / verdict logic (v0.5.0b1 alpha stubs) ───────────────────────────

def _evaluate_claim(claim_data: dict) -> tuple[str, list[str]]:
    """v0.5.0b1 stub: minimal clearance check. v1.0: full JisPolicy + Rust verify."""
    denials: list[str] = []
    actor = claim_data.get("actor", "")
    if not actor:
        denials.append("missing_actor")
    clearance = claim_data.get("clearance", 0)
    if not isinstance(clearance, int) or clearance < 1:
        denials.append("clearance_too_low")
    verdict = "denied" if denials else "allowed"
    return verdict, denials


def _next_fira_phase(phase: str) -> tuple[float, Optional[str]]:
    """First-Initiation Revoke/Accept state machine (alpha stub).

    F = First-init     → trust 0.5, awaits A or R
    A = Accept         → trust 0.7, terminal
    R = Revoke         → trust 0.0, terminal
    """
    if phase == "F":
        return 0.5, "A"     # next step is Accept (or implicit R)
    if phase == "A":
        return 0.7, None
    if phase == "R":
        return 0.0, None
    raise ValueError(f"invalid FIR/A phase: {phase!r}")


# ── Local fallbacks for soft-bootstrap (TIBET_SOFT_BOOTSTRAP=1) ─────────────

def _local_claim(sess: "Session", claim_data: dict) -> dict:
    verdict, denials = _evaluate_claim(claim_data)
    return {"kind": "RESULT", "ok": True, "verdict": verdict,
            "denials": denials, "_local": True}


def _local_bind(sess: "Session", actor: str, intent: str, ttl_seconds: int) -> dict:
    binding_token = f"bind_{secrets.token_hex(8)}"
    expires_at = time.time() + ttl_seconds
    sess._local_state.setdefault("bindings", {})[binding_token] = {
        "actor": actor, "intent": intent, "expires_at": expires_at,
    }
    return {"kind": "RESULT", "ok": True, "binding_token": binding_token,
            "expires_at": expires_at, "_local": True}


def _local_fira(sess: "Session", peer: str, phase: str) -> dict:
    trust, next_phase = _next_fira_phase(phase)
    sess._local_state.setdefault("fira", {})[peer] = {"trust": trust, "phase": phase}
    return {"kind": "RESULT", "ok": True, "trust": trust,
            "next_phase": next_phase, "_local": True}


# ── Soft-bootstrap (TIBET_SOFT_BOOTSTRAP=1, shared env-var) ─────────────────

def _soft_bootstrap(actor: str, reason: str) -> Session:
    print(
        f"[jis-core] WARNING: TIBET_SOFT_BOOTSTRAP active — degraded to ephemeral identity. "
        f"actor={actor!r} reason={reason!r}. NOT bound to central JIS store. "
        f"This is for dev/test only; production must run with a live jis-OSAPI.",
        file=sys.stderr,
    )
    return Session(
        actor=actor,
        session_token=_new_session_token(),
        chain_handle="local:ephemeral",
        identity_binding="local:ephemeral",
        is_soft_bootstrap=True,
    )


# ── Client: bootstrap() ─────────────────────────────────────────────────────

def bootstrap(
    actor: str,
    actor_claim: bytes,
    *,
    url: Optional[str] = None,
    capabilities: Optional[list[str]] = None,
    timeout: float = 5.0,
) -> Session:
    """Perform the three-message jis-OSAPI handshake (spec §3) and return a Session.

    For a non-kernel package the canonical pattern is to call BOTH:
        tibet_session = tibet_core.bootstrap(actor=..., actor_claim=...)
        jis_session   = jis_core.bootstrap(actor=..., actor_claim=...)
    The handshake is identical; the operations differ (provenance vs identity).
    """
    if not isinstance(actor_claim, (bytes, bytearray)):
        raise IdentityError("actor_claim must be bytes (JIS-signed)")
    if not actor_claim:
        raise IdentityError("actor_claim is empty")

    resolved = _resolve_url(url)
    soft = os.environ.get("TIBET_SOFT_BOOTSTRAP") == "1"

    try:
        sock = _connect(resolved, timeout=timeout)
    except (OSError, socket.timeout) as e:
        if soft:
            return _soft_bootstrap(actor, reason=f"connect-failed: {e}")
        raise BootstrapError(f"cannot reach jis-OSAPI at {resolved}: {e}") from None

    try:
        _send(sock, {
            "kind": "HELLO",
            "v": PROTOCOL_VERSION,
            "actor": actor,
            "actor_claim": actor_claim.hex(),
            "capabilities": capabilities or ["claim", "bind"],
            "pid": os.getpid(),
            "encoding": "ldjson",
        })

        msg = _recv_line(sock)
        kind = msg.get("kind")
        if kind == "ERROR":
            code = msg.get("code", "PROTOCOL")
            detail = msg.get("detail", "")
            raise {
                "IDENTITY": IdentityError,
                "VERSION": VersionError,
            }.get(code, ProtocolError)(f"jis-OSAPI rejected handshake: {code} — {detail}")
        if kind != "WELCOME":
            raise ProtocolError(f"expected WELCOME, got {kind!r}")

        version = msg.get("v")
        if version not in SUPPORTED_VERSIONS:
            raise VersionError(f"jis-OSAPI advertised v={version!r}, client supports {SUPPORTED_VERSIONS}")

        session_token = msg["session_token"]

        _send(sock, {
            "kind": "ACK",
            "session_token": session_token,
            "echo": _ack_hmac(session_token),
        })

        return Session(
            actor=actor,
            session_token=session_token,
            chain_handle=msg["chain_handle"],
            identity_binding=msg["identity_binding"],
            heartbeat_interval_ms=msg.get("heartbeat_interval_ms", DEFAULT_HEARTBEAT_INTERVAL_MS),
            soft_stop_topic=msg.get("soft_stop_topic", SOFT_STOP_TOPIC),
            server_url=resolved,
        )
    finally:
        try:
            sock.close()
        except Exception:
            pass


# ── Server (reference implementation; daemon-thread based) ──────────────────

ClaimVerifier = Callable[[str, bytes], bool]
"""(actor, claim_bytes) -> ok. v1.0 wires this to the real Ed25519 verify via jis_core.jis_core."""


def _default_claim_verifier(actor: str, claim: bytes) -> bool:
    """v0.5.0 stub: accept any non-empty claim. v1.0: jis_core.verify_claim_signature()."""
    return len(claim) > 0


@dataclass
class OSAPIServer:
    """Reference jis-OSAPI server — minimal, threaded, for tests + dev.

    Production: the real jis-OSAPI daemon (separate package, v0.6.x).
    """

    bind_url: str = field(default_factory=lambda: f"tcp://{DEFAULT_TCP_HOST}:0")
    claim_verifier: ClaimVerifier = field(default=_default_claim_verifier)
    _sock: Optional[socket.socket] = None
    _thread: Optional[threading.Thread] = None
    _stop: threading.Event = field(default_factory=threading.Event)
    _bound_url: str = ""
    _active_sessions: dict[str, dict] = field(default_factory=dict)

    @property
    def url(self) -> str:
        return self._bound_url or self.bind_url

    def start(self) -> None:
        url = self.bind_url
        if url.startswith("unix://"):
            path = url[len("unix://"):]
            try:
                os.unlink(path)
            except FileNotFoundError:
                pass
            os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
            s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            s.bind(path)
            self._bound_url = f"unix://{path}"
        else:
            if url.startswith("tcp://"):
                url = url[len("tcp://"):]
            host, _, port = url.partition(":")
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            s.bind((host, int(port or 0)))
            actual_port = s.getsockname()[1]
            self._bound_url = f"tcp://{host}:{actual_port}"

        s.listen(8)
        s.settimeout(0.5)
        self._sock = s
        self._thread = threading.Thread(target=self._accept_loop, daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        if self._sock is not None:
            try:
                self._sock.close()
            except Exception:
                pass
        if self._thread is not None:
            self._thread.join(timeout=2.0)

    def _accept_loop(self) -> None:
        assert self._sock is not None
        while not self._stop.is_set():
            try:
                conn, _addr = self._sock.accept()
            except (socket.timeout, OSError):
                continue
            threading.Thread(target=self._handle, args=(conn,), daemon=True).start()

    def _handle(self, conn: socket.socket) -> None:
        """Dispatcher: HELLO → handshake, OP → operation, anything-else → ERROR."""
        try:
            conn.settimeout(5.0)
            try:
                first = _recv_line(conn)
            except ProtocolError as e:
                _send(conn, {"kind": "ERROR", "code": "PROTOCOL", "detail": str(e)})
                return

            kind = first.get("kind")
            if kind == "HELLO":
                self._handle_handshake(conn, first)
            elif kind == "OP":
                self._handle_op(conn, first)
            else:
                _send(conn, {"kind": "ERROR", "code": "PROTOCOL",
                             "detail": f"expected HELLO or OP, got {kind!r}"})
        finally:
            try:
                conn.close()
            except Exception:
                pass

    def _handle_handshake(self, conn: socket.socket, hello: dict) -> None:
        client_v = hello.get("v")
        if client_v not in SUPPORTED_VERSIONS:
            _send(conn, {"kind": "ERROR", "code": "VERSION",
                          "detail": f"server supports {list(SUPPORTED_VERSIONS)}, client offered {client_v!r}"})
            return

        actor = hello.get("actor") or ""
        claim_hex = hello.get("actor_claim") or ""
        try:
            claim = bytes.fromhex(claim_hex)
        except ValueError:
            _send(conn, {"kind": "ERROR", "code": "IDENTITY", "detail": "actor_claim not valid hex"})
            return

        if not actor or not self.claim_verifier(actor, claim):
            _send(conn, {"kind": "ERROR", "code": "IDENTITY", "detail": "claim verification failed"})
            return

        session_token = _new_session_token()
        chain_handle = f"ch_{secrets.token_hex(8)}"
        identity_binding = f"id_{secrets.token_hex(8)}"

        _send(conn, {
            "kind": "WELCOME",
            "v": PROTOCOL_VERSION,
            "session_token": session_token,
            "chain_handle": chain_handle,
            "identity_binding": identity_binding,
            "heartbeat_interval_ms": DEFAULT_HEARTBEAT_INTERVAL_MS,
            "soft_stop_topic": SOFT_STOP_TOPIC,
            "supported_versions": list(SUPPORTED_VERSIONS),
        })

        try:
            ack = _recv_line(conn)
        except ProtocolError:
            return

        expected_echo = _ack_hmac(session_token)
        if ack.get("kind") != "ACK" or ack.get("echo") != expected_echo:
            return

        self._active_sessions[session_token] = {
            "actor": actor,
            "chain_handle": chain_handle,
            "identity_binding": identity_binding,
            "started_at": time.time(),
            "bindings": {},      # binding_token -> {actor, intent, expires_at}
            "fira": {},          # peer -> {trust, phase}
        }

    def _handle_op(self, conn: socket.socket, msg: dict) -> None:
        token = msg.get("session_token")
        sess = self._active_sessions.get(token) if token else None
        if not sess:
            _send(conn, {"kind": "ERROR", "code": "SESSION",
                         "detail": "unknown or expired session_token"})
            return

        op = msg.get("op")
        if op == "claim":
            self._op_claim(conn, sess, msg)
        elif op == "bind":
            self._op_bind(conn, sess, msg)
        elif op == "fira":
            self._op_fira(conn, sess, msg)
        else:
            _send(conn, {"kind": "ERROR", "code": "OP_UNKNOWN",
                         "detail": f"unknown op: {op!r}; supported: claim/bind/fira"})

    # ── Op handlers (spec §4 jis-OSAPI) ──────────────────────────────────────

    def _op_claim(self, conn: socket.socket, sess: dict, msg: dict) -> None:
        claim_data = msg.get("claim") or {}
        if not isinstance(claim_data, dict):
            _send(conn, {"kind": "ERROR", "code": "OP_INVALID",
                         "detail": "claim must be an object"})
            return
        verdict, denials = _evaluate_claim(claim_data)
        _send(conn, {"kind": "RESULT", "ok": True,
                     "verdict": verdict, "denials": denials})

    def _op_bind(self, conn: socket.socket, sess: dict, msg: dict) -> None:
        actor = msg.get("actor") or ""
        intent = msg.get("intent") or ""
        ttl = int(msg.get("ttl_seconds", 3600))
        if not actor or not intent:
            _send(conn, {"kind": "ERROR", "code": "OP_INVALID",
                         "detail": "bind requires actor + intent"})
            return
        binding_token = f"bind_{secrets.token_hex(8)}"
        expires_at = time.time() + ttl
        sess["bindings"][binding_token] = {"actor": actor, "intent": intent, "expires_at": expires_at}
        _send(conn, {"kind": "RESULT", "ok": True,
                     "binding_token": binding_token, "expires_at": expires_at})

    def _op_fira(self, conn: socket.socket, sess: dict, msg: dict) -> None:
        peer = msg.get("peer") or ""
        phase = msg.get("phase") or ""
        if not peer or phase not in {"F", "R", "A"}:
            _send(conn, {"kind": "ERROR", "code": "OP_INVALID",
                         "detail": "fira requires peer + phase ∈ {F,R,A}"})
            return
        try:
            trust, next_phase = _next_fira_phase(phase)
        except ValueError as e:
            _send(conn, {"kind": "ERROR", "code": "OP_INVALID", "detail": str(e)})
            return
        sess["fira"][peer] = {"trust": trust, "phase": phase}
        _send(conn, {"kind": "RESULT", "ok": True,
                     "trust": trust, "next_phase": next_phase})
