#!/usr/bin/env python3
"""Codex review MCP server for ARIS, backed by the Codex app-server protocol.

Replaces the `codex mcp-server` subcommand that Codex CLI removed in 0.155.0.
Codex still ships an app-server (JSON-RPC over stdio); this bridge speaks MCP to
the harness on one side and the app-server protocol on the other, exposing the
same two tools the old MCP server did:

    codex        start a Codex session, run one turn, return the reply
    codex-reply  continue an existing thread by id

ARIS reads `structuredContent.threadId` to continue a review thread, so every
successful call returns it alongside the text content.

The app-server child is spawned once and kept alive for the process lifetime:
threads live inside it, and a restart would make previously returned threadIds
unresumable. `thread/resume` is used for replies whose thread this process has
not seen, which also covers a thread started before a bridge restart.
"""

from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
import threading
import time
import traceback
import uuid
from pathlib import Path
from typing import Any

# ── stdio framing ───────────────────────────────────────────────────────────


def _configure_stdio_for_mcp() -> None:
    """Re-bind stdio as raw binary so the MCP framing is never re-encoded."""
    sys.stdout = os.fdopen(sys.stdout.fileno(), "wb", buffering=0)
    sys.stdin = os.fdopen(sys.stdin.fileno(), "rb", buffering=0)


SERVER_NAME = os.environ.get("CODEX_REVIEW_SERVER_NAME", "codex")
TOOL_CALL_TIMEOUT_SEC = int(os.environ.get("CODEX_REVIEW_TIMEOUT_SEC", "1200000")) / 1000.0
_default_debug = os.environ.get("TEMP", "/tmp") if sys.platform == "win32" else "/tmp"
DEBUG_LOG = Path(os.environ.get("CODEX_REVIEW_DEBUG_LOG", f"{_default_debug}/{SERVER_NAME}-mcp-debug.log"))

_use_ndjson = False


def debug_log(message: str) -> None:
    if not os.environ.get("CODEX_REVIEW_DEBUG"):
        return
    try:
        DEBUG_LOG.parent.mkdir(parents=True, exist_ok=True)
        with DEBUG_LOG.open("a", encoding="utf-8") as fh:
            fh.write(f"{message}\n")
    except OSError:
        pass


def send_response(response: dict[str, Any]) -> None:
    payload = json.dumps(response, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    debug_log(f"SEND {payload.decode('utf-8', errors='replace')[:2000]}")
    if _use_ndjson:
        sys.stdout.write(payload + b"\n")
    else:
        sys.stdout.write(f"Content-Length: {len(payload)}\r\n\r\n".encode("utf-8") + payload)
    sys.stdout.flush()


def read_message() -> dict[str, Any] | None:
    global _use_ndjson

    line = sys.stdin.readline()
    if not line:
        return None

    text = line.decode("utf-8").rstrip("\r\n")
    if text.lower().startswith("content-length:"):
        try:
            length = int(text.split(":", 1)[1].strip())
        except ValueError:
            return None
        while True:
            header = sys.stdin.readline()
            if not header or header in (b"\r\n", b"\n"):
                break
        body = sys.stdin.read(length)
        if not body:
            return None
        try:
            return json.loads(body.decode("utf-8"))
        except json.JSONDecodeError:
            return None

    _use_ndjson = True
    if not text:
        return read_message()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return None


# ── codex binary resolution ─────────────────────────────────────────────────


def resolve_codex_bin() -> list[str]:
    """Return argv prefix that runs the Codex CLI.

    On Windows the npm shim is a .cmd, which CreateProcess cannot execute
    directly, so prefer the vendored executable the shim ultimately launches.
    """
    override = os.environ.get("CODEX_BIN")
    if override:
        return [override]

    npm_root = Path(os.environ.get("APPDATA", "")) / "npm" / "node_modules" / "@openai" / "codex"
    vendor = npm_root / "node_modules" / "@openai"
    if vendor.is_dir():
        for candidate in sorted(vendor.glob("codex-*/vendor/*/bin/codex.exe")):
            if candidate.is_file():
                return [str(candidate)]

    found = shutil.which("codex")
    if found:
        return [found]
    raise FileNotFoundError("codex CLI not found; set CODEX_BIN")


# ── app-server client ───────────────────────────────────────────────────────

INIT_TIMEOUT_SEC = 60.0


class CodexAppServer:
    """One long-lived `codex app-server` child, multiplexed by request id."""

    def __init__(self, cwd: str) -> None:
        self.cwd = cwd
        self._proc: subprocess.Popen[bytes] | None = None
        self._lock = threading.Lock()
        self._send_lock = threading.Lock()
        self._next_id = 0
        self._pending: dict[int, dict[str, Any]] = {}
        self._events: dict[str, list[dict[str, Any]]] = {}
        self._done: dict[str, threading.Event] = {}
        self._reader: threading.Thread | None = None
        # Threads this process started or resumed. A threadId in here needs no
        # `thread/resume` before the next turn.
        self._known: set[str] = set()

    # -- lifecycle

    def _ensure_started(self) -> None:
        with self._lock:
            if self._proc is not None and self._proc.poll() is None:
                return
            argv = resolve_codex_bin() + ["app-server"]
            debug_log(f"SPAWN {' '.join(argv)} (cwd={self.cwd})")
            self._proc = subprocess.Popen(
                argv,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.DEVNULL,
                cwd=self.cwd,
                bufsize=0,
            )
            self._pending.clear()
            self._reader = threading.Thread(target=self._read_loop, daemon=True)
            self._reader.start()
            self._call("initialize", {
                "clientInfo": {"name": SERVER_NAME, "version": "1.0.0"},
                "capabilities": {},
            }, timeout=INIT_TIMEOUT_SEC)

    def _read_loop(self) -> None:
        proc = self._proc
        if proc is None or proc.stdout is None:
            return
        for raw in proc.stdout:
            line = raw.decode("utf-8", "replace").strip()
            if not line:
                continue
            try:
                msg = json.loads(line)
            except json.JSONDecodeError:
                continue
            debug_log(f"RECV {line[:1500]}")

            rid = msg.get("id")
            if rid is not None and "method" not in msg:
                with self._send_lock:
                    waiter = self._pending.pop(rid, None)
                if waiter is not None:
                    waiter["response"] = msg
                    waiter["event"].set()
                continue

            params = msg.get("params") or {}
            thread_id = params.get("threadId")
            if isinstance(thread_id, str):
                with self._send_lock:
                    bucket = self._events.setdefault(thread_id, [])
                    bucket.append(msg)
                    done = self._done.get(thread_id)
                if done is not None and msg.get("method") == "turn/completed":
                    done.set()

    def _write(self, payload: dict[str, Any]) -> None:
        proc = self._proc
        if proc is None or proc.stdin is None:
            raise RuntimeError("codex app-server is not running")
        proc.stdin.write((json.dumps(payload) + "\n").encode("utf-8"))
        proc.stdin.flush()

    def _call(self, method: str, params: dict[str, Any], timeout: float = INIT_TIMEOUT_SEC) -> dict[str, Any]:
        with self._send_lock:
            self._next_id += 1
            rid = self._next_id
            waiter = {"event": threading.Event(), "response": None}
            self._pending[rid] = waiter
        self._write({"jsonrpc": "2.0", "id": rid, "method": method, "params": params})
        if not waiter["event"].wait(timeout):
            raise TimeoutError(f"codex app-server did not answer {method} within {timeout:.0f}s")
        response = waiter["response"] or {}
        if "error" in response:
            raise RuntimeError(f"{method} failed: {json.dumps(response['error'])[:400]}")
        return response.get("result") or {}

    # -- threads and turns

    def start_thread(self, params: dict[str, Any]) -> str:
        self._ensure_started()
        opts = dict(params)
        opts.setdefault("cwd", self.cwd)
        result = self._call("thread/start", opts)
        thread = result.get("thread") or {}
        thread_id = thread.get("id") or result.get("threadId")
        if not isinstance(thread_id, str) or not thread_id:
            raise RuntimeError(f"thread/start returned no thread id: {json.dumps(result)[:300]}")
        with self._send_lock:
            self._known.add(thread_id)
        return thread_id

    def has_thread(self, thread_id: str) -> bool:
        """Whether this process already holds the thread open in its app-server."""
        with self._send_lock:
            return thread_id in self._known

    def resume_thread(self, thread_id: str, params: dict[str, Any]) -> None:
        self._ensure_started()
        self._call("thread/resume", {"threadId": thread_id, **params})
        with self._send_lock:
            self._known.add(thread_id)

    def run_turn(self, thread_id: str, prompt: str, params: dict[str, Any]) -> str:
        with self._send_lock:
            self._events[thread_id] = []
            self._done[thread_id] = threading.Event()
            done = self._done[thread_id]

        turn_params: dict[str, Any] = {"threadId": thread_id, "input": [{"type": "text", "text": prompt}]}
        for key in ("model", "effort", "approvalPolicy", "sandboxPolicy", "cwd", "outputSchema"):
            if params.get(key) is not None:
                turn_params[key] = params[key]
        self._call("turn/start", turn_params, timeout=60.0)

        if not done.wait(TOOL_CALL_TIMEOUT_SEC):
            raise TimeoutError(f"codex turn did not complete within {TOOL_CALL_TIMEOUT_SEC:.0f}s")

        with self._send_lock:
            events = list(self._events.pop(thread_id, []))
            self._done.pop(thread_id, None)

        # Prefer the finalized turn payload; fall back to the last agent message.
        for msg in reversed(events):
            if msg.get("method") != "turn/completed":
                continue
            turn = (msg.get("params") or {}).get("turn") or {}
            if turn.get("status") == "failed":
                raise RuntimeError(f"codex turn failed: {json.dumps(turn.get('error'))[:400]}")
            text = _agent_text_from_items(turn.get("items"))
            if text is not None:
                return text

        texts = []
        for msg in events:
            if msg.get("method") != "item/completed":
                continue
            item = (msg.get("params") or {}).get("item") or {}
            if item.get("type") == "agentMessage" and isinstance(item.get("text"), str):
                texts.append(item["text"])
        if texts:
            return "\n\n".join(texts)
        raise RuntimeError("codex turn completed without an agent message")


def _agent_text_from_items(items: Any) -> str | None:
    if not isinstance(items, list):
        return None
    texts = [i.get("text") for i in items
             if isinstance(i, dict) and i.get("type") == "agentMessage" and isinstance(i.get("text"), str)]
    return "\n\n".join(texts) if texts else None


# ── tool schemas (the contract the old codex mcp-server exposed) ────────────

CODEX_TOOL = {
    "name": "codex",
    "description": "Run a Codex session. Accepts configuration parameters matching the Codex Config struct.",
    "inputSchema": {
        "type": "object",
        "properties": {
            "prompt": {"type": "string", "description": "The prompt to send to Codex."},
            "model": {"type": "string", "description": "Model to use (e.g. gpt-5.1-codex)."},
            "cwd": {"type": "string", "description": "Working directory for the session."},
            "sandbox": {"type": "string", "description": "Sandbox mode: read-only, workspace-write, danger-full-access."},
            "approval-policy": {"type": "string", "description": "Approval policy: untrusted, on-failure, on-request, never."},
            "config": {"type": "object", "description": "Additional Codex configuration overrides."},
            "base-instructions": {"type": "string", "description": "Override the base system instructions."},
            "developer-instructions": {"type": "string", "description": "Developer instructions appended to the system prompt."},
            "compact-prompt": {"type": "string", "description": "Prompt used when compacting the conversation."},
        },
        "required": ["prompt"],
    },
}

CODEX_REPLY_TOOL = {
    "name": "codex-reply",
    "description": "Continue a Codex conversation by providing the thread id and prompt.",
    "inputSchema": {
        "type": "object",
        "properties": {
            "prompt": {"type": "string", "description": "The follow-up prompt to send to Codex."},
            "threadId": {"type": "string", "description": "Thread id returned by a previous codex call."},
            "conversationId": {"type": "string", "description": "Alias for threadId."},
        },
        "required": ["prompt"],
    },
}


def _thread_params(args: dict[str, Any]) -> dict[str, Any]:
    params: dict[str, Any] = {}
    if isinstance(args.get("model"), str):
        params["model"] = args["model"]
    if isinstance(args.get("cwd"), str):
        params["cwd"] = args["cwd"]
    if isinstance(args.get("sandbox"), str):
        params["sandbox"] = args["sandbox"]
    if isinstance(args.get("approval-policy"), str):
        params["approvalPolicy"] = args["approval-policy"]
    if isinstance(args.get("base-instructions"), str):
        params["baseInstructions"] = args["base-instructions"]
    if isinstance(args.get("developer-instructions"), str):
        params["developerInstructions"] = args["developer-instructions"]
    if isinstance(args.get("config"), dict):
        params["config"] = args["config"]
    return params


def _result(text: str, thread_id: str) -> dict[str, Any]:
    return {
        "content": [{"type": "text", "text": text}],
        "structuredContent": {"threadId": thread_id},
        "isError": False,
    }


def _error(text: str) -> dict[str, Any]:
    return {"content": [{"type": "text", "text": text}], "isError": True}


def handle_tool_call(server: CodexAppServer, name: str, args: dict[str, Any]) -> dict[str, Any]:
    prompt = args.get("prompt")
    if not isinstance(prompt, str) or not prompt:
        return _error("prompt is required")

    try:
        if name == "codex":
            params = _thread_params(args)
            thread_id = server.start_thread(params)
            text = server.run_turn(thread_id, prompt, params)
            return _result(text, thread_id)

        if name == "codex-reply":
            thread_id = args.get("threadId") or args.get("conversationId")
            if not isinstance(thread_id, str) or not thread_id:
                return _error("threadId is required for codex-reply")
            if not server.has_thread(thread_id):
                server.resume_thread(thread_id, _thread_params(args))
            text = server.run_turn(thread_id, prompt, _thread_params(args))
            return _result(text, thread_id)

        return _error(f"unknown tool: {name}")
    except Exception as exc:  # surface the reason to the reviewer, never crash the bridge
        debug_log(f"TOOL ERROR {name}: {traceback.format_exc()}")
        return _error(f"codex bridge: {type(exc).__name__}: {exc}")


# ── MCP protocol ────────────────────────────────────────────────────────────


def main() -> int:
    _configure_stdio_for_mcp()
    cwd = os.environ.get("CODEX_REVIEW_CWD") or str(Path.home())
    server = CodexAppServer(cwd)

    while True:
        message = read_message()
        if message is None:
            break

        method = message.get("method")
        rid = message.get("id")

        if method == "initialize":
            send_response({
                "jsonrpc": "2.0",
                "id": rid,
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {"tools": {"listChanged": False}},
                    "serverInfo": {"name": SERVER_NAME, "version": "1.0.0"},
                },
            })
        elif method == "tools/list":
            send_response({"jsonrpc": "2.0", "id": rid, "result": {"tools": [CODEX_TOOL, CODEX_REPLY_TOOL]}})
        elif method == "tools/call":
            params = message.get("params") or {}
            name = params.get("name")
            args = params.get("arguments") or {}
            result = handle_tool_call(server, str(name), args if isinstance(args, dict) else {})
            send_response({"jsonrpc": "2.0", "id": rid, "result": result})
        elif method in ("notifications/initialized", "notifications/cancelled"):
            pass
        elif rid is not None:
            send_response({
                "jsonrpc": "2.0",
                "id": rid,
                "error": {"code": -32601, "message": f"method not found: {method}"},
            })

    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(0)