"""Localhost WebSocket server: token handshake, engine detection, run streaming."""

from __future__ import annotations

import asyncio
import contextlib
import json
from collections.abc import Awaitable, Callable

import websockets

from atashi_bridge import pairing
from atashi_bridge.engine_auth import detect_engine_auth
from atashi_bridge.engines import detect_engines
from atashi_bridge.events import BridgeEvent
from atashi_bridge.runner import run_engine


class BridgeServer:
    def __init__(
        self,
        pairing_token: str,
        engine_detector: Callable[[], list[str]] = detect_engines,
        engine_auth_detector: Callable[[list[str]], Awaitable[dict]] = detect_engine_auth,
        approver: Callable[[str, str], Awaitable[bool]] = pairing.approve,
        roots_desc: str = "",
    ):
        self._token = pairing_token
        self._detect = engine_detector
        self._detect_auth = engine_auth_detector
        self._approve = approver
        self._roots_desc = roots_desc

    async def _handle(self, ws):
        try:
            first = json.loads(await ws.recv())
        except (json.JSONDecodeError, websockets.ConnectionClosed):
            await ws.close()
            return

        mtype = first.get("type")
        if mtype == "hello":
            if first.get("token") != self._token:
                await ws.close(code=4401, reason="unauthorized")
                return
            await ws.send(json.dumps({"type": "ready"}))
        elif mtype == "pair_request":
            origin = first.get("origin", "unknown")
            # The approver can block for up to its timeout (default 60s) waiting
            # for a human at the terminal; the client may give up and disconnect
            # in that window, so guard the reply against a closed socket.
            approved = await self._approve(origin, self._roots_desc)
            try:
                if not approved:
                    print(f"Pairing denied for origin: {origin}", flush=True)
                    await ws.send(json.dumps({"type": "error", "code": "pairing_denied"}))
                    await ws.close(code=4403, reason="pairing denied")
                    return
                print(f"Pairing approved for origin: {origin}", flush=True)
                await ws.send(json.dumps({"type": "paired", "token": self._token}))
            except websockets.ConnectionClosed:
                return
        else:
            await ws.close(code=4401, reason="unauthorized")
            return

        # Protocol v2 (seamless talk/work): runs execute as a task so the
        # receive loop keeps reading — a cancel message arriving mid-run can
        # actually be processed. Under the old inline-await structure a cancel
        # could never be read. One active run per connection; sends are
        # serialised so the run task and this loop never interleave frames.
        run_task: asyncio.Task | None = None
        send_lock = asyncio.Lock()

        async def _send(payload: dict) -> None:
            async with send_lock:
                await ws.send(json.dumps(payload))

        try:
            async for raw in ws:
                try:
                    msg = json.loads(raw)
                except json.JSONDecodeError:
                    continue
                mtype = msg.get("type")
                if mtype == "detect":
                    # Signed-in state alongside installed state (2026-08-10,
                    # talk-workspace-context spec item 2): detect_engines()
                    # is shutil.which() only, so an engine can be offered
                    # that fails on the first message because the local CLI
                    # was never logged in. engine_auth values are
                    # True/False/None (codex, claude_code) or always None
                    # (antigravity — no known status command).
                    detected = self._detect()
                    auth = await self._detect_auth(detected)
                    await _send({"type": "engines", "engines": detected, "engine_auth": auth})
                elif mtype == "run":
                    if run_task and not run_task.done():
                        await _send(BridgeEvent.error("a run is already active on this connection").to_dict())
                        continue
                    run_task = asyncio.create_task(self._run(_send, msg))
                elif mtype == "cancel":
                    if run_task and not run_task.done():
                        run_task.cancel()
                        with contextlib.suppress(asyncio.CancelledError):
                            await run_task
                        with contextlib.suppress(websockets.ConnectionClosed):
                            await _send(BridgeEvent.cancelled().to_dict())
        finally:
            # Client went away mid-run: stop the engine, don't orphan it.
            if run_task and not run_task.done():
                run_task.cancel()
                with contextlib.suppress(asyncio.CancelledError):
                    await run_task

    async def _run(self, send, msg):
        try:
            engine = msg["engine"]
            prompt = msg["prompt"]
            workspace_dir = msg["workspace_dir"]
        except KeyError as exc:
            await send(BridgeEvent.error(f"missing field: {exc}").to_dict())
            return
        try:
            async for event in run_engine(
                engine=engine,
                prompt=prompt,
                model=msg.get("model"),
                workspace_dir=workspace_dir,
                api_key=msg.get("api_key"),
                kind=msg.get("kind") or "run",
                resume_session_id=msg.get("resume_session_id"),
                mcp_servers=msg.get("mcp_servers"),
                extra_env=msg.get("env"),
                cold_start_reason=msg.get("cold_start_reason"),
                agent_identity=msg.get("agent_identity"),
            ):
                await send(event.to_dict())
        except ValueError as exc:
            await send(BridgeEvent.error(str(exc)).to_dict())

    @contextlib.asynccontextmanager
    async def serving(self, host: str = "127.0.0.1", port: int = 8765):
        async with websockets.serve(self._handle, host, port) as server:
            bound = server.sockets[0].getsockname()[1]
            yield bound
