"""JSON-RPC 2.0 over stdio client for the MemOS bridge.

Spawns the packaged bridge subprocess (compiled ``dist/bridge.cjs`` when
available, otherwise the source ``bridge.cts`` through ``tsx``) and communicates
via line-delimited JSON messages on its stdin/stdout. Responses are matched by
``id``. Notifications (events + logs) are forwarded to registered callbacks on a
reader thread.

The client is *blocking* by design — callers wanting async behaviour
should wrap requests in a thread pool.
"""

from __future__ import annotations

import contextlib
import hashlib
import json
import logging
import os
import queue
import shutil
import subprocess
import threading

from pathlib import Path
from typing import TYPE_CHECKING, Any

from runtime_home import resolve_runtime_home


if TYPE_CHECKING:
    from collections.abc import Callable


logger = logging.getLogger(__name__)

HOST_HANDLER_WAIT_SECONDS = 5.0
HOST_HANDLER_QUEUE_CAPACITY = 16

# ─── Module-level singleton tracker ─────────────────────────────────────
# Each entry maps an ``(agent, no_viewer, runtime_home)`` key to the
# most-recent active ``MemosBridgeClient`` for that slot. When a new client
# is constructed for an existing key, the previous client is closed
# synchronously so the Node-side ``bridge.cjs`` subprocess does not leak.
#
# This is the Python-side guard against issue #1910 (bridge process leak:
# every turn spawns new bridge.cjs). Defence in depth on the Node side
# lives in ``bridge.cts`` via ``bridge-stdio.pid``.
_ACTIVE_CLIENTS: dict[tuple[str, bool, str], MemosBridgeClient] = {}
_ACTIVE_CLIENTS_LOCK = threading.Lock()


def _expanded_path(value: str, env: dict[str, str]) -> Path:
    """Resolve a path using the child process' HOME rather than global state."""
    home = env.get("HOME", "").strip() or str(Path.home())
    if value == "~":
        value = home
    elif value.startswith(("~/", "~\\")):
        value = str(Path(home) / value[2:])
    return Path(value).resolve()


def _resolved_runtime_home(agent: str, env: dict[str, str]) -> Path:
    """Mirror the Node home resolver for singleton/process ownership."""
    memos_home = env.get("MEMOS_HOME", "").strip()
    if memos_home:
        return _expanded_path(memos_home, env)
    config_file = env.get("MEMOS_CONFIG_FILE", "").strip()
    if config_file:
        return _expanded_path(config_file, env).parent
    if agent == "hermes":
        plugin_root = Path(__file__).resolve().parent.parent.parent.parent
        return resolve_runtime_home(env=env, plugin_root=plugin_root)
    agent_home = f".{agent}"
    default_home = Path(env.get("HOME", "") or Path.home()) / agent_home / "memos-plugin"
    return _expanded_path(str(default_home), env)


def _runtime_scope_token(agent: str, runtime_home: Path) -> str:
    """Return a stable, path-private token safe for use in a PID filename."""
    raw = f"{agent}\0{runtime_home}".encode()
    return hashlib.sha256(raw).hexdigest()[:24]


def _installed_node_binary(plugin_root: Path) -> str | None:
    marker = plugin_root / ".memos-node-bin"
    try:
        candidate = marker.read_text(encoding="utf-8").strip()
    except OSError:
        return None
    if candidate and os.path.isfile(candidate) and os.access(candidate, os.X_OK):
        return candidate
    return None


def _bridge_script(plugin_root: Path) -> Path:
    """Pick the bridge entrypoint, preferring pure ESM over the CJS trampoline.

    Resolution order (issue #1736):
        1. ``dist/bridge.mjs`` — pure ESM compiled output, the only entry
           that avoids the CJS↔ESM bridge that fails on Node ≥ 22.
        2. ``dist/bridge.cjs`` — legacy CommonJS compiled output, kept for
           installations whose ``dist/`` predates the ESM entrypoint.
        3. ``bridge.mts`` — pure ESM TypeScript source for ``tsx``-driven
           local development.
        4. ``bridge.cts`` — legacy CommonJS TypeScript source. Returned
           as the last-resort default so error messages stay stable when
           none of the candidates exist.
    """
    candidates = (
        plugin_root / "dist" / "bridge.mjs",
        plugin_root / "dist" / "bridge.cjs",
        plugin_root / "bridge.mts",
    )
    for candidate in candidates:
        if candidate.exists():
            return candidate
    return plugin_root / "bridge.cts"


class BridgeError(RuntimeError):
    """Raised when the bridge returns a JSON-RPC error object."""

    def __init__(self, code: str, message: str, data: Any = None) -> None:
        super().__init__(f"[{code}] {message}")
        self.code = code
        self.message = message
        self.data = data


class MemosBridgeClient:
    """Client wrapping a line-delimited JSON-RPC 2.0 stdio bridge.

    Usage:
        >>> client = MemosBridgeClient()
        >>> client.request("core.health", {})
        {'ok': True, 'version': '...'}
        >>> client.close()

    Thread-safe: per-request locking ensures concurrent callers don't
    interleave writes.
    """

    def __init__(
        self,
        *,
        bridge_path: str | None = None,
        node_binary: str | None = None,
        agent: str = "hermes",
        no_viewer: bool = True,
        extra_env: dict[str, str] | None = None,
        runtime_home: str | None = None,
    ) -> None:
        self._lock = threading.Lock()
        self._next_id = 1
        self._pending: dict[int, dict[str, Any]] = {}
        self._events: list[Callable[[dict[str, Any]], None]] = []
        self._logs: list[Callable[[dict[str, Any]], None]] = []
        # Reverse-direction handlers: the bridge can send us a
        # JSON-RPC request via `serverRequest(...)` (e.g.
        # `host.llm.complete` for fallback LLM calls). Registered
        # methods run on one bounded, daemon worker. Keeping execution
        # serial preserves the adapter's previous concurrency contract while
        # preventing a slow host LLM call from blocking stdout response
        # demultiplexing for every shared provider lease.
        self._host_handlers: dict[str, Callable[[dict[str, Any]], Any]] = {}
        self._host_handlers_cv = threading.Condition()
        self._host_handler_queue: queue.Queue[tuple[Any, str, dict[str, Any]] | None] = queue.Queue(
            maxsize=HOST_HANDLER_QUEUE_CAPACITY
        )
        self._host_handler_stop = threading.Event()
        self._closed = False

        plugin_root = Path(__file__).resolve().parent.parent.parent.parent
        node = (
            node_binary
            or os.environ.get("MEMOS_NODE_BINARY")
            or _installed_node_binary(plugin_root)
            or shutil.which("node")
            or "node"
        )
        script_path = Path(bridge_path) if bridge_path else _bridge_script(plugin_root)
        script = str(script_path)
        env = {**os.environ, **(extra_env or {})}
        resolved_runtime_home = (
            _expanded_path(runtime_home, env)
            if runtime_home
            else _resolved_runtime_home(agent, env)
        )
        if runtime_home and not {
            "MEMOS_HOME",
            "MEMOS_CONFIG_FILE",
        }.intersection(extra_env or {}):
            # An explicit home is both an ownership scope and the child
            # runtime location. Callers that deliberately snapshot either
            # selector in extra_env retain that more specific configuration.
            env["MEMOS_HOME"] = str(resolved_runtime_home)
            env.pop("MEMOS_CONFIG_FILE", None)
        runtime_scope = _runtime_scope_token(agent, resolved_runtime_home)

        # Prefer the compiled JavaScript bridge — the new pure ESM
        # ``dist/bridge.mjs`` (issue #1736) or the legacy ``dist/bridge.cjs``
        # — both run on plain ``node`` without any loader. The raw
        # TypeScript entries remain as a development fallback and need
        # ``tsx`` for stripping types plus ``.js`` → ``.ts`` import
        # resolution. On Windows the ``.bin/tsx`` file is a shell shim,
        # so use tsx's real JS entrypoint whenever we have to launch the
        # source entry through a specific Node.
        tsx_cli = plugin_root / "node_modules" / "tsx" / "dist" / "cli.mjs"
        bridge_args = [
            script,
            f"--agent={agent}",
            f"--runtime-scope={runtime_scope}",
        ]
        if no_viewer:
            bridge_args.append("--no-viewer")
        if script_path.suffix in (".mjs", ".cjs"):
            cmd = [node, *bridge_args]
        elif tsx_cli.exists():
            cmd = [node, str(tsx_cli), *bridge_args]
        else:
            # Fallback path: `node --import tsx` reproduces the same loader
            # inline. Requires tsx to be resolvable as a package from the
            # plugin root — true whenever node_modules exists. If tsx is
            # genuinely missing the child will fail fast with a loader
            # error the stderr reader will surface.
            cmd = [node, "--import", "tsx", *bridge_args]
        self._proc = subprocess.Popen(
            cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding="utf-8",
            errors="replace",
            bufsize=1,
            env=env,
            cwd=str(plugin_root),
        )
        self._host_handler_worker = threading.Thread(
            target=self._host_handler_loop,
            daemon=True,
            name="memos-bridge-host-handler",
        )
        self._host_handler_worker.start()
        self._reader = threading.Thread(
            target=self._read_loop,
            daemon=True,
            name="memos-bridge-reader",
        )
        self._reader.start()
        self._stderr_reader = threading.Thread(
            target=self._stderr_loop,
            daemon=True,
            name="memos-bridge-stderr",
        )
        self._stderr_reader.start()

        # Singleton tracking (issue #1910). Register ourselves as the
        # active client for ``(agent, no_viewer, runtime_home)`` and reap
        # any previous holder synchronously so its subprocess does not leak.
        # The reap happens AFTER our reader threads are running, so the
        # previous client's ``close()`` (which closes stdin and waits for
        # exit) cannot interfere with our own startup.
        self._singleton_agent = agent
        self._singleton_no_viewer = bool(no_viewer)
        self._singleton_runtime_home = str(resolved_runtime_home)
        previous = self._register_active()
        if previous is not None and previous is not self:
            prev_pid = getattr(previous, "pid", "?")
            logger.info(
                "MemOS: closing previous bridge client (pid=%s) before adopting new one (pid=%s)",
                prev_pid,
                self.pid,
            )
            with contextlib.suppress(Exception):
                previous.close()

    def _register_active(self) -> MemosBridgeClient | None:
        """Register self as the active singleton; return the displaced client."""
        key = (
            self._singleton_agent,
            self._singleton_no_viewer,
            self._singleton_runtime_home,
        )
        with _ACTIVE_CLIENTS_LOCK:
            previous = _ACTIVE_CLIENTS.get(key)
            _ACTIVE_CLIENTS[key] = self
        return previous

    def _unregister_active(self) -> None:
        """Remove self from the active registry if we are still the current entry."""
        key = (
            self._singleton_agent,
            self._singleton_no_viewer,
            self._singleton_runtime_home,
        )
        with _ACTIVE_CLIENTS_LOCK:
            if _ACTIVE_CLIENTS.get(key) is self:
                _ACTIVE_CLIENTS.pop(key, None)

    @property
    def pid(self) -> int:
        """Return the PID of the bridge subprocess."""
        return int(getattr(self._proc, "pid", 0) or 0)

    # ─── Public API ──

    def request(
        self,
        method: str,
        params: Any = None,
        *,
        timeout: float = 30.0,
    ) -> dict[str, Any]:
        if self._closed:
            raise BridgeError("transport_closed", "bridge client is closed")
        # Layer 2 (#2028): if the subprocess has already exited, bail
        # BEFORE writing to the pipe. Otherwise the write silently
        # buffers into a dead pipe and the caller parks on the full
        # per-request timeout.
        exit_code = None
        with contextlib.suppress(Exception):
            exit_code = self._proc.poll()
        if exit_code is not None:
            raise BridgeError(
                "transport_closed",
                f"bridge subprocess exited (code={exit_code})",
            )
        with self._lock:
            rpc_id = self._next_id
            self._next_id += 1
            waiter = threading.Event()
            entry: dict[str, Any] = {"event": waiter, "result": None, "error": None}
            self._pending[rpc_id] = entry
            payload = json.dumps(
                {"jsonrpc": "2.0", "id": rpc_id, "method": method, "params": params},
                ensure_ascii=False,
            )
            try:
                self._proc.stdin.write(payload + "\n")
                self._proc.stdin.flush()
            except (BrokenPipeError, OSError) as err:
                self._pending.pop(rpc_id, None)
                raise BridgeError("transport_closed", str(err)) from err

        if not waiter.wait(timeout=timeout):
            with self._lock:
                self._pending.pop(rpc_id, None)
            raise BridgeError("timeout", f"{method} did not respond within {timeout}s")
        if entry["error"] is not None:
            e = entry["error"]
            raise BridgeError(
                e.get("data", {}).get("code") or str(e.get("code", "internal")),
                e.get("message", "unknown error"),
                e.get("data"),
            )
        return entry["result"] or {}

    def notify(self, method: str, params: Any = None) -> None:
        if self._closed:
            return
        with self._lock:
            payload = json.dumps({"jsonrpc": "2.0", "method": method, "params": params})
            try:
                self._proc.stdin.write(payload + "\n")
                self._proc.stdin.flush()
            except (BrokenPipeError, OSError, ValueError):
                pass

    def on_event(self, cb: Callable[[dict[str, Any]], None]) -> None:
        self._events.append(cb)

    def on_log(self, cb: Callable[[dict[str, Any]], None]) -> None:
        self._logs.append(cb)

    def register_host_handler(
        self,
        method: str,
        handler: Callable[[dict[str, Any]], Any],
    ) -> None:
        """Register a handler for bridge → adapter (reverse) requests.

        The Node-side bridge calls these via ``stdio.serverRequest``.
        Most-recent registration wins. Handlers run serially on a bounded
        daemon worker so a long-running host LLM call cannot stall the reader
        thread that resolves unrelated foreground JSON-RPC responses.
        """
        with self._host_handlers_cv:
            self._host_handlers[method] = handler
            self._host_handlers_cv.notify_all()

    def close(self) -> None:
        if self._closed:
            return
        with self._host_handlers_cv:
            self._closed = True
            self._host_handlers_cv.notify_all()
        self._stop_host_handler_worker()

        # Drop self from the module-level singleton tracker (issue #1910)
        # BEFORE the potentially-slow stdin/SIGTERM/SIGKILL dance. We
        # only evict the registry slot if we still own it — a newer
        # client that displaced us must remain reachable.
        self._unregister_active()

        pid = self.pid

        # 1. Close stdin (triggers bridge's graceful exit)
        with contextlib.suppress(Exception):
            self._proc.stdin.close()

        # 2. Wait for process to exit gracefully (up to 5 seconds)
        try:
            self._proc.wait(timeout=5.0)
            logger.debug("MemOS: bridge process %d exited gracefully", pid)
        except subprocess.TimeoutExpired:
            # 3. If still running, send SIGTERM
            logger.warning(
                "MemOS: bridge process %d did not exit after stdin close, sending SIGTERM", pid
            )
            try:
                self._proc.terminate()  # Send SIGTERM
                self._proc.wait(timeout=5.0)  # Increased from 2.0 to 5.0 for viewer cleanup
                logger.debug("MemOS: bridge process %d terminated", pid)
            except subprocess.TimeoutExpired:
                # 4. Last resort: SIGKILL
                logger.error(
                    "MemOS: bridge process %d did not respond to SIGTERM, sending SIGKILL", pid
                )
                self._proc.kill()  # Send SIGKILL
                try:
                    self._proc.wait(timeout=1.0)
                except subprocess.TimeoutExpired:
                    logger.error("MemOS: bridge process %d could not be killed", pid)

        # 5. Clean up pending requests
        self._abort_pending("bridge closed")

    # ─── Internals ──

    def _abort_pending(self, reason: str) -> None:
        """Wake every parked JSON-RPC waiter with `transport_closed`.

        Called from both `close()` and the `_read_loop` `finally:`
        block, so any pending request sees a real transport error
        immediately instead of parking on its per-request timeout.
        Also flips `_closed` and notifies host-handler waiters so a
        second `request()` fails fast.
        """
        with self._host_handlers_cv:
            self._closed = True
            self._host_handlers_cv.notify_all()
        self._stop_host_handler_worker()
        with self._lock:
            for entry in list(self._pending.values()):
                entry["error"] = {
                    "code": -32000,
                    "message": reason,
                    "data": {"code": "transport_closed"},
                }
                entry["event"].set()
            self._pending.clear()

    def _read_loop(self) -> None:
        assert self._proc.stdout is not None
        try:
            for line in self._proc.stdout:
                line = line.strip()
                if not line:
                    continue
                try:
                    msg = json.loads(line)
                except json.JSONDecodeError:
                    logger.debug("bridge: malformed line: %r", line[:120])
                    continue
                if "id" in msg and msg["id"] is not None and ("result" in msg or "error" in msg):
                    self._resolve(msg)
                    continue
                if msg.get("method") == "events.notify":
                    for cb in list(self._events):
                        try:
                            cb(msg.get("params") or {})
                        except Exception:
                            logger.debug("event listener threw", exc_info=True)
                    continue
                if msg.get("method") == "logs.forward":
                    for cb in list(self._logs):
                        try:
                            cb(msg.get("params") or {})
                        except Exception:
                            logger.debug("log listener threw", exc_info=True)
                    continue
                # Reverse-direction request: the bridge is asking the
                # adapter to do something (e.g. run a fallback LLM call
                # via `host.llm.complete`). Dispatch to the registered
                # handler on the bounded worker. The reader must return to
                # stdout immediately so a slow host LLM callback cannot
                # head-of-line block normal JSON-RPC responses.
                method = msg.get("method")
                rpc_id = msg.get("id")
                if (
                    isinstance(method, str)
                    and rpc_id is not None
                    and "result" not in msg
                    and "error" not in msg
                ):
                    params = msg.get("params") or {}
                    if not isinstance(params, dict):
                        params = {}
                    self._dispatch_host_request(rpc_id, method, params)
                    continue
        except Exception:
            # Any unexpected exception in the reader loop still needs
            # to fall through to the `finally:` cleanup so callers
            # don't park on 30 s timeouts.
            logger.debug("bridge reader thread crashed", exc_info=True)
        finally:
            # Layer 1 (#2028): the reader thread is the only signal
            # that the Node subprocess has stopped answering. On any
            # exit — normal EOF or exception — abort pending waiters
            # immediately so callers get `transport_closed` in < 1 s
            # instead of waiting for each 30 s per-request timeout.
            self._abort_pending("bridge subprocess exited")

    def _dispatch_host_request(
        self,
        rpc_id: Any,
        method: str,
        params: dict[str, Any],
    ) -> None:
        """Queue reverse RPC work without ever blocking the reader thread."""
        if self._closed:
            return
        try:
            self._host_handler_queue.put_nowait((rpc_id, method, params))
        except queue.Full:
            logger.warning("host handler queue full; rejecting %s", method)
            self._send_response(
                rpc_id,
                error={
                    "code": -32000,
                    "message": "host handler queue is full",
                    "data": {"code": "host_handler_busy"},
                },
            )

    def _host_handler_loop(self) -> None:
        """Run reverse RPC handlers serially away from stdout demultiplexing."""
        while True:
            request = self._host_handler_queue.get()
            try:
                if request is None:
                    return
                rpc_id, method, params = request
                if self._closed:
                    continue
                handler = self._host_handler_for(method)
                if handler is None:
                    self._send_response(
                        rpc_id,
                        error={
                            "code": -32601,
                            "message": f"method not found: {method}",
                            "data": {"code": "unknown_method"},
                        },
                    )
                    continue
                try:
                    result = handler(params)
                    self._send_response(rpc_id, result=result)
                except Exception as err:
                    logger.warning("host handler %s failed: %s", method, err)
                    self._send_response(
                        rpc_id,
                        error={
                            "code": -32000,
                            "message": str(err) or err.__class__.__name__,
                            "data": {"code": "host_handler_failed"},
                        },
                    )
            finally:
                self._host_handler_queue.task_done()

    def _stop_host_handler_worker(self) -> None:
        """Discard queued callbacks and ask the daemon worker to exit."""
        if self._host_handler_stop.is_set():
            return
        self._host_handler_stop.set()
        while True:
            try:
                self._host_handler_queue.get_nowait()
            except queue.Empty:
                break
            else:
                self._host_handler_queue.task_done()
        with contextlib.suppress(queue.Full):
            self._host_handler_queue.put_nowait(None)

    def _host_handler_for(
        self,
        method: str,
        *,
        timeout: float = HOST_HANDLER_WAIT_SECONDS,
    ) -> Callable[[dict[str, Any]], Any] | None:
        """Return a reverse-RPC handler, waiting briefly during startup.

        The Node bridge now starts stdio before ``core.init()`` so host LLM
        fallback can run during startup recovery. On a fast machine that
        reverse request can arrive just before ``initialize()`` registers
        ``host.llm.complete``. Waiting here turns that sub-millisecond race
        into the intended handshake while still returning ``unknown_method``
        for genuinely unsupported methods.
        """
        with self._host_handlers_cv:
            self._host_handlers_cv.wait_for(
                lambda: method in self._host_handlers or self._closed,
                timeout=timeout,
            )
            return self._host_handlers.get(method)

    def _send_response(
        self,
        rpc_id: Any,
        *,
        result: Any = None,
        error: dict[str, Any] | None = None,
    ) -> None:
        """Write a JSON-RPC response for a reverse-direction request."""
        if self._closed:
            return
        payload: dict[str, Any] = {"jsonrpc": "2.0", "id": rpc_id}
        if error is not None:
            payload["error"] = error
        else:
            payload["result"] = result
        with self._lock:
            try:
                self._proc.stdin.write(json.dumps(payload, ensure_ascii=False) + "\n")
                self._proc.stdin.flush()
            except (BrokenPipeError, OSError, ValueError):
                pass

    def _stderr_loop(self) -> None:
        assert self._proc.stderr is not None
        for line in self._proc.stderr:
            line = line.rstrip()
            if line:
                logger.debug("bridge.stderr: %s", line)

    def _resolve(self, msg: dict[str, Any]) -> None:
        rpc_id = msg.get("id")
        if not isinstance(rpc_id, int):
            return
        with self._lock:
            entry = self._pending.pop(rpc_id, None)
        if not entry:
            return
        if "error" in msg:
            entry["error"] = msg["error"]
        else:
            entry["result"] = msg.get("result")
        entry["event"].set()