"""Minimal Codex app-server stdio JSON-RPC client for bootstrap.

Bootstrap runs on the engineer host. It launches a `codex app-server` subprocess
and speaks newline-delimited JSON (JSONL) over stdio.

The protocol is JSON-RPC 2.0, but the `"jsonrpc":"2.0"` header is omitted by
codex-app-server and should be omitted by clients as well.

We intentionally keep this implementation small and dependency-free so
`flywheel-bootstrap` remains easy to run via uvx in BYOC environments.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Mapping
import json
import queue
import subprocess
import threading

JsonDict = dict[str, Any]


class AppServerProtocolError(RuntimeError):
    pass


class AppServerRequestError(RuntimeError):
    def __init__(self, *, code: int | None, message: str, data: object | None) -> None:
        self.code = code
        self.message = message
        self.data = data
        super().__init__(message)


@dataclass(frozen=True)
class AppServerInitializeResponse:
    user_agent: str


def _jsonrpc_request(*, req_id: int, method: str, params: JsonDict) -> JsonDict:
    return {"id": req_id, "method": method, "params": params}


def _jsonrpc_notification(*, method: str, params: JsonDict | None = None) -> JsonDict:
    if params is None:
        return {"method": method}
    return {"method": method, "params": params}


def _jsonrpc_response(*, req_id: object, result: JsonDict) -> JsonDict:
    return {"id": req_id, "result": result}


def _jsonrpc_error(*, req_id: object, code: int, message: str) -> JsonDict:
    return {"id": req_id, "error": {"code": code, "message": message}}


class AppServerStdioClient:
    """Threaded JSON-RPC client for `codex app-server` over stdio."""

    def __init__(
        self,
        *,
        argv: list[str],
        env: Mapping[str, str] | None,
        cwd: Path,
        on_notification: Callable[[JsonDict], None] | None = None,
        request_handler: Callable[[JsonDict], JsonDict] | None = None,
    ) -> None:
        self._argv = list(argv)
        self._env = dict(env) if env is not None else None
        self._cwd = cwd
        self._on_notification = on_notification
        self._request_handler = request_handler

        self._proc: subprocess.Popen[str] | None = None
        self._reader_thread: threading.Thread | None = None
        self._send_lock = threading.Lock()
        self._pending_lock = threading.Lock()
        self._pending: dict[object, queue.Queue[JsonDict]] = {}
        self._closed = threading.Event()

        self._next_id_lock = threading.Lock()
        self._next_id = 1

    def start(self) -> None:
        if self._proc is not None:
            return
        proc = subprocess.Popen(
            self._argv,
            cwd=self._cwd,
            env=self._env,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            # IMPORTANT: do not pipe stderr without draining it, otherwise
            # a chatty server can fill the OS pipe buffer and deadlock.
            stderr=None,
            text=True,
            encoding="utf-8",
            bufsize=1,
        )
        if proc.stdin is None or proc.stdout is None:
            proc.kill()
            raise AppServerProtocolError("Failed to open stdio pipes for app-server")
        self._proc = proc
        self._reader_thread = threading.Thread(target=self._reader_loop, daemon=True)
        self._reader_thread.start()

    def close(self) -> None:
        self._closed.set()
        proc = self._proc
        if proc is None:
            return
        try:
            proc.terminate()
            proc.wait(timeout=2)
        except Exception:
            pass
        self._proc = None

    def _alloc_id(self) -> int:
        with self._next_id_lock:
            req_id = self._next_id
            self._next_id += 1
        return req_id

    def request(
        self, *, method: str, params: JsonDict, timeout_seconds: float = 60
    ) -> JsonDict:
        self.start()
        req_id = self._alloc_id()
        q: queue.Queue[JsonDict] = queue.Queue(maxsize=1)
        with self._pending_lock:
            self._pending[req_id] = q
        self._send(_jsonrpc_request(req_id=req_id, method=method, params=params))
        try:
            msg = q.get(timeout=timeout_seconds)
        finally:
            with self._pending_lock:
                self._pending.pop(req_id, None)

        if "error" in msg:
            err = msg.get("error")
            if isinstance(err, dict):
                code = err.get("code")
                message = err.get("message")
                data = err.get("data")
                raise AppServerRequestError(
                    code=int(code) if isinstance(code, int) else None,
                    message=str(message)
                    if isinstance(message, str) and message
                    else "app-server error",
                    data=data,
                )
            raise AppServerRequestError(code=None, message="app-server error", data=err)
        result = msg.get("result")
        if not isinstance(result, dict):
            raise AppServerProtocolError(f"Malformed JSON-RPC response for {method}")
        return result

    def notify(self, *, method: str, params: JsonDict | None = None) -> None:
        self.start()
        self._send(_jsonrpc_notification(method=method, params=params))

    def initialize(
        self,
        *,
        client_info: JsonDict,
        experimental_api: bool = False,
    ) -> AppServerInitializeResponse:
        params: JsonDict = {
            "clientInfo": dict(client_info),
            "capabilities": {"experimentalApi": bool(experimental_api)},
        }
        result = self.request(method="initialize", params=params, timeout_seconds=60)
        user_agent = result.get("userAgent")
        if not isinstance(user_agent, str) or not user_agent:
            raise AppServerProtocolError("initialize missing userAgent")

        # Required handshake ack: clients must emit `initialized` notification.
        self.notify(method="initialized")
        return AppServerInitializeResponse(user_agent=user_agent)

    def _send(self, msg: JsonDict) -> None:
        proc = self._proc
        if proc is None or proc.stdin is None:
            raise AppServerProtocolError("app-server subprocess not running")
        text = json.dumps(msg, separators=(",", ":")) + "\n"
        with self._send_lock:
            try:
                proc.stdin.write(text)
                proc.stdin.flush()
            except BrokenPipeError as exc:
                raise AppServerProtocolError("app-server subprocess exited") from exc

    def _handle_server_request(self, msg: JsonDict) -> None:
        req_id = msg.get("id")
        handler = self._request_handler
        if handler is None:
            self._send(
                _jsonrpc_error(req_id=req_id, code=-32601, message="not supported")
            )
            return
        try:
            result = handler(dict(msg))
            if not isinstance(result, dict):
                raise AppServerProtocolError(
                    "request_handler must return a dict result"
                )
            self._send(_jsonrpc_response(req_id=req_id, result=result))
        except Exception as exc:  # pragma: no cover - defensive
            self._send(_jsonrpc_error(req_id=req_id, code=-32000, message=str(exc)))

    def _reader_loop(self) -> None:
        proc = self._proc
        if proc is None or proc.stdout is None:
            return
        for line in proc.stdout:
            if self._closed.is_set():
                break
            raw = line.strip()
            if not raw:
                continue
            try:
                parsed = json.loads(raw)
            except json.JSONDecodeError:
                continue
            if not isinstance(parsed, dict):
                continue

            # Server-initiated request: has id+method but no result/error.
            if (
                "id" in parsed
                and "method" in parsed
                and "result" not in parsed
                and "error" not in parsed
            ):
                self._handle_server_request(parsed)
                continue

            # Response to a prior request.
            if "id" in parsed and ("result" in parsed or "error" in parsed):
                req_id = parsed.get("id")
                q: queue.Queue[JsonDict] | None = None
                with self._pending_lock:
                    q = self._pending.get(req_id)
                if q is not None:
                    try:
                        q.put_nowait(parsed)
                    except queue.Full:
                        pass
                continue

            # Notification.
            if "method" in parsed and "id" not in parsed:
                cb = self._on_notification
                if cb is not None:
                    try:
                        cb(dict(parsed))
                    except Exception:
                        # Avoid killing the reader thread for UI mapping bugs.
                        pass
                continue
