"""
CodexCLIAdapter — drives OpenAI's `codex` CLI via per-turn subprocess.

Each turn spawns:
  codex exec --json [--sandbox workspace-write] <text>

Session history is maintained by codex itself in ~/.codex/history.jsonl.
To continue the last session: codex resume --last --json <text>

JSONL event stream (stdout):
  {type: "thread.started",  thread: {id}}
  {type: "turn.started"}
  {type: "item.started",    item: {id, type, ...}}
  {type: "item.delta",      item: {id, type, delta}}   ← streaming text
  {type: "item.completed",  item: {id, type, ...}}
  {type: "turn.completed",  usage: {input_tokens, output_tokens, ...}}
  {type: "turn.failed",     error: {message}}

Item types in item.started/item.completed:
  agent_message      — TextBlock (content field)
  command_execution  — ToolCallBlock (command field)
  mcp_tool_call      — ToolCallBlock (tool / params fields)
  file_change        — ToolCallBlock (path, change_type fields)
  web_search         — ToolCallBlock (query field)
  reasoning          — ThinkingBlock (summary field)
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import shutil
import uuid
from pathlib import Path
from typing import AsyncIterator, Callable

from ..attachments import Attachment, build_text_with_inline
from ..blocks import (
    Block, AgentInfoBlock, AudioHintBlock, ErrorBlock, SessionEndBlock,
    SessionStatusHintBlock, TextBlock, ThinkingBlock, ToolCallBlock, tool_result_or_large,
)
from .base import PerTurnCLIAdapter, ExternalSession, active_cwds_for_command
from ..sentinel_tools import codex_mcp_overrides

log = logging.getLogger("beacon.codex")

# Codex sandbox levels: read-only | workspace-write | danger-full-access
_DEFAULT_SANDBOX = "workspace-write"

# Codex reasoning-effort levels (config key `model_reasoning_effort`). GPT-5-family models
# accept minimal|low|medium|high; surfaced as an effort picker like claude's --effort.
_CODEX_EFFORT_LEVELS = ["minimal", "low", "medium", "high"]

# Mirror the other per-turn adapters' active-turn window: transcript writes within this many
# seconds count as an in-progress turn; anything older is a quiet-but-open session.
_TURN_RECENT_SECS = 20.0

# Curated fallback model catalog. Codex has NO `models list` CLI command and its `-m` flag
# accepts any id (validated server-side), so this is a hand-kept preset list — it does NOT
# auto-refresh; edit it + restart the bridge to change it. list_models() also folds in
# whatever `model` the user actually configured in ~/.codex/config.toml so their real model
# always appears (and is preselected) even when it's newer than this list.
_CODEX_STATIC_MODELS = [
    {"id": "gpt-5-codex", "label": "GPT-5 Codex", "detail": "Codex-tuned"},
    {"id": "gpt-5",       "label": "GPT-5",       "detail": "General reasoning"},
    {"id": "gpt-5-mini",  "label": "GPT-5 mini",  "detail": "Faster, cheaper"},
    {"id": "o3",          "label": "o3",          "detail": "Reasoning"},
    {"id": "o4-mini",     "label": "o4-mini",     "detail": "Fast reasoning"},
]


def _codex_config_models() -> list[str]:
    """Model ids the user actually configured in ~/.codex/config.toml (top-level `model`
    plus any `[profiles.*].model`). Lets the picker track their real setup without a code
    edit — the one part of the catalog that isn't hardcoded (still read once, at startup)."""
    cfg = Path.home() / ".codex" / "config.toml"
    if not cfg.exists():
        return []
    try:
        import tomllib
        data = tomllib.loads(cfg.read_text())
    except Exception:
        return []
    out: list[str] = []
    top = data.get("model")
    if isinstance(top, str) and top.strip():
        out.append(top.strip())
    profiles = data.get("profiles")
    if isinstance(profiles, dict):
        for p in profiles.values():
            m = p.get("model") if isinstance(p, dict) else None
            if isinstance(m, str) and m.strip() and m.strip() not in out:
                out.append(m.strip())
    return out


_CODEX_SESSIONS = Path.home() / ".codex" / "sessions"


def _codex_rollout_meta(path: Path) -> tuple[str, str, str]:
    """(session_id, cwd, first_user_prompt) from a codex rollout's head."""
    sid = cwd = prompt = ""
    try:
        with path.open("r", errors="replace") as f:
            for _ in range(80):
                line = f.readline()
                if not line:
                    break
                try:
                    o = json.loads(line)
                except Exception:
                    continue
                t = o.get("type"); p = o.get("payload") or {}
                if t == "session_meta" and not sid:
                    sid = p.get("id", "") or ""
                elif t == "turn_context" and not cwd:
                    cwd = p.get("cwd", "") or ""
                elif t == "response_item" and not prompt and p.get("type") == "message" and p.get("role") == "user":
                    for c in (p.get("content") or []):
                        if isinstance(c, dict) and c.get("type") == "input_text" and c.get("text"):
                            tx = c["text"].strip()
                            if not tx.startswith("<"):
                                prompt = tx[:200]
                            break
                if sid and cwd and prompt:
                    break
    except Exception:
        pass
    return sid, cwd, prompt


def _codex_rollouts_by_mtime() -> list[tuple[float, Path]]:
    if not _CODEX_SESSIONS.exists():
        return []
    out: list[tuple[float, Path]] = []
    for f in _CODEX_SESSIONS.rglob("rollout-*.jsonl"):
        try:
            out.append((f.stat().st_mtime, f))
        except OSError:
            continue
    out.sort(key=lambda x: x[0], reverse=True)
    return out


def _scan_codex_sessions(limit: int = 20) -> list[ExternalSession]:
    import time as _t
    owned = {os.path.realpath(c) for c in active_cwds_for_command("/codex")}
    now = _t.time()
    out: list[ExternalSession] = []
    for mtime, f in _codex_rollouts_by_mtime()[:limit]:
        sid, cwd, prompt = _codex_rollout_meta(f)
        if not sid:
            continue
        is_active = bool(cwd) and os.path.realpath(cwd) in owned
        activity = ("turn" if (now - mtime) <= _TURN_RECENT_SECS else "open") if is_active else "idle"
        out.append(ExternalSession(session_id=sid, agent_id="codex", cwd=cwd,
                                   last_prompt=prompt, started_at="", is_active=is_active,
                                   activity=activity))
    return out


class CodexCLIAdapter(PerTurnCLIAdapter):
    cli_name = "codex"
    _capture_stderr = False   # stderr → DEVNULL (codex is noisy; the tail is rc-based)
    _close_stdin = True       # codex appends stdin to the prompt → close it or it hangs
    cap_images = True
    cap_documents = True  # PDFs/doc attachments inlined as text via build_text_with_inline
    # P4: codex DOES load the beacon MCP, but in non-interactive `exec` its sandbox cancels
    # the MCP tool call ("user cancelled MCP tool call") and sanitizes the child's env. We
    # deliver the tools by (a) injecting the port+session into the MCP child via per-spawn
    # `-c mcp_servers.beacon.env.*` (codex_mcp_overrides) and (b) spawning with
    # --dangerously-bypass-approvals-and-sandbox so the call is allowed. NOTE: the bypass
    # also disables shell-command sandboxing for codex turns — accepted by default here.
    cap_sentinel = True
    cap_thinking = True
    cap_interrupt = True
    cap_input_during_turn = False  # serializes sends in its own _send_queue (no true
                                   # mid-turn interleave) — route mid-turn msgs through
                                   # the bridge's droppable queue, not an invisible defer
    provides_models = True         # -m <model>; catalog = static presets + ~/.codex config
    provides_usage = True
    agent_id = "codex"
    display_name = "Codex CLI"
    beta = False              # live-verified via the adapter matrix
    cli_binary = "codex"

    def __init__(
        self,
        binary: str | None = None,
        cwd: str | None = None,
        sandbox: str = _DEFAULT_SANDBOX,
        resume_last: bool = True,
        resume_session_id: str | None = None,
        on_session_id: "Callable[[str], None] | None" = None,
        model: str | None = None,
        effort: str | None = None,
    ) -> None:
        self._binary = binary or shutil.which("codex") or "codex"
        self._cwd = cwd
        self._sandbox = sandbox
        self._resume_last = resume_last
        # Per-turn model (`-m`) + reasoning effort (`-c model_reasoning_effort`). None ⇒ codex
        # falls back to ~/.codex/config.toml defaults. Re-applied on the NEXT turn's spawn.
        self._model = (model or "").strip() or None
        self._effort = (effort or "").strip() or None
        # Codex thread/session UUID — resume by id (turns after the first, or an adopted
        # external session). Captured from the thread.started event. docs/session-handoff.md
        self._thread_id: str | None = resume_session_id
        self._on_session_id = on_session_id
        self._is_first_turn = True
        self._send_queue: asyncio.Queue[tuple[str, list[Attachment]]] = None  # type: ignore
        self._pending_items: dict[str, dict] = {}   # item_id → item dict from .started
        self._fetch_store: dict[str, str] = {}
        self._proc: asyncio.subprocess.Process | None = None
        self._usage: dict = {}

    # ── Session handoff (docs/session-handoff.md) ──────────────────────────────

    @classmethod
    def scan_external_sessions(cls) -> list[ExternalSession]:
        return _scan_codex_sessions()

    @classmethod
    def transcript_path(cls, ext_sid: str) -> "Path | None":
        """Find the rollout JSONL whose parsed session id matches `ext_sid`."""
        for _mtime, f in _codex_rollouts_by_mtime():
            try:
                sid, _, _ = _codex_rollout_meta(f)
                if sid == ext_sid:
                    return f
            except Exception:
                pass
        return None

    @classmethod
    def has_changes_since(cls, ts: float) -> bool:
        """True if any codex session file has been written since `ts`."""
        if not _CODEX_SESSIONS.exists():
            return False
        try:
            if _CODEX_SESSIONS.stat().st_mtime > ts:
                return True
        except OSError:
            pass
        for f in _CODEX_SESSIONS.rglob("rollout-*.jsonl"):
            try:
                if f.stat().st_mtime > ts:
                    return True
            except OSError:
                pass
        return False

    @classmethod
    def is_cwd_externally_owned(cls, cwd: str, sid: str = "") -> bool:
        if not cwd:
            return False
        target = os.path.realpath(os.path.expanduser(cwd))
        return any(os.path.realpath(c) == target for c in active_cwds_for_command("/codex"))

    @classmethod
    def latest_sid_for_cwd(cls, cwd: str) -> str | None:
        if not cwd:
            return None
        target = os.path.realpath(os.path.expanduser(cwd))
        for _mtime, f in _codex_rollouts_by_mtime():   # newest first
            sid, c, _ = _codex_rollout_meta(f)
            if sid and c and os.path.realpath(c) == target:
                return sid
        return None

    @classmethod
    def make_transcript_mirror(cls, cwd: str, sid: str, *, from_end: bool = False):
        from ..transcript_mirror import make_codex_mirror
        return make_codex_mirror(sid, from_end=from_end)

    @classmethod
    def terminal_resume_command(cls, sid: str) -> str | None:
        import shlex
        return f"codex resume {shlex.quote(sid)}" if sid else "codex"

    @classmethod
    def readiness(cls) -> tuple[bool, str]:
        if (Path.home() / ".codex" / "auth.json").exists() or os.environ.get("OPENAI_API_KEY"):
            return True, ""
        return False, "Run `codex login` (or set OPENAI_API_KEY)."

    @classmethod
    async def list_models(cls, extra: dict | None = None) -> list[dict]:
        # Static curated presets + whatever the user configured in ~/.codex/config.toml (so a
        # model newer than the presets — e.g. gpt-5.4-mini — still shows up, preselected).
        # Codex exposes no model-list endpoint, so this only refreshes at bridge startup.
        out: list[dict] = []
        seen: set[str] = set()
        for mid in _codex_config_models():
            out.append({"id": mid, "label": mid, "detail": "From ~/.codex/config.toml"})
            seen.add(mid)
        for m in _CODEX_STATIC_MODELS:
            if m["id"] not in seen:
                out.append(m)
                seen.add(m["id"])
        return out

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    async def start(self, session_id: str) -> AgentInfoBlock:
        self._beacon_session_id = session_id   # stamped into per-turn spawn env (beacon MCP routing)
        self._send_queue = asyncio.Queue()
        self._is_first_turn = True
        return self.make_agent_info(
            name="Codex CLI",
            working_dirs=[self._cwd] if self._cwd else [],
            effort_levels=_CODEX_EFFORT_LEVELS,
            effort=self._effort or "",
        )

    def set_model(self, model: str) -> None:
        """Per-turn adapter: stored, applied on the NEXT turn's spawn (`-m`)."""
        self._model = (model or "").strip() or None

    def set_effort(self, level: str) -> None:
        """Per-turn adapter: stored, applied on the NEXT turn's spawn
        (`-c model_reasoning_effort`)."""
        self._effort = (level or "").strip() or None

    def clear_history(self) -> None:
        self._is_first_turn = True
        self._thread_id = None
        # Wipe the persisted ext_sid too, else a reconnect respawns with --resume and
        # resurrects the old context. Next turn spawns fresh; codex mints a new thread id.
        if self._on_session_id:
            self._on_session_id("")

    # ── Turn loop (spawn/stream/teardown handled by PerTurnCLIAdapter) ─────────

    async def _build_argv(self, text: str, attachments: list[Attachment]) -> list[str]:
        # Resume the specific session by id once known (turn ≥2, or an adopted external
        # session); else a fresh `codex exec`. Resume-by-id beats `--last` — it can't
        # cross-talk with another codex session that became most-recent meanwhile.
        # --skip-git-repo-check: the bridge drives codex in arbitrary cwds, many not git
        # repos; without it codex refuses to run outside a trusted/git dir.
        # `codex exec` is non-interactive (never prompts) and, as of codex 0.139, rejects
        # `--ask-for-approval`. We spawn with --dangerously-bypass-approvals-and-sandbox
        # (replaces --sandbox): non-interactive codex otherwise CANCELS MCP tool calls, so
        # the beacon sentinel tools never fire. Both `exec` and `exec resume` accept it and
        # the `-c` sentinel overrides (which wire the beacon MCP child to this session).
        sentinel_ov = (codex_mcp_overrides(self._beacon_session_id)
                       if type(self).cap_sentinel else [])
        # Model (`-m`) + reasoning effort (`-c model_reasoning_effort="…"`, TOML-quoted like
        # the sentinel overrides). Both `exec` and `exec resume` accept them; unset ⇒ codex
        # uses its ~/.codex/config.toml defaults.
        model_ov: list[str] = []
        if self._model:
            model_ov += ["-m", self._model]
        if self._effort:
            model_ov += ["-c", f'model_reasoning_effort="{self._effort}"']
        bypass = "--dangerously-bypass-approvals-and-sandbox"
        # `codex exec` takes the prompt as a single positional arg; images can be attached
        # via --image, while everything else is inlined as text/path hints. See attachments.py.
        image_args: list[str] = []
        for att in attachments:
            if att.is_image:
                image_args += ["--image", str(att.local_path)]
        preamble = self.sentinel_client_preamble(is_first_turn=not self._thread_id)
        prompt = preamble + await asyncio.to_thread(build_text_with_inline, text, attachments)
        self._is_first_turn = False
        if self._thread_id:
            return [self._binary, "exec", "resume", "--json", bypass, "--skip-git-repo-check",
                    *model_ov, *sentinel_ov, *image_args, self._thread_id, prompt]
        return [self._binary, "exec", "--json", bypass,
                "--skip-git-repo-check", *model_ov, *sentinel_ov, *image_args, prompt]

    def _reset_turn_state(self) -> None:
        self._pending_items.clear()

    async def get_usage(self) -> dict:
        return dict(self._usage)

    # ── Event parser ──────────────────────────────────────────────────────────

    def _todo_card(self, item: dict) -> Block:
        """Map a codex `todo_list` item → the generic live TodoBlock.

        codex's `update_plan` surfaces as a `todo_list` item carrying
        `items:[{text, completed}]` (re-sent in full on every change). There's no
        in_progress state — a step is either pending or completed.
        """
        todos = [
            {"content": t.get("text", ""),
             "status": "completed" if t.get("completed") else "pending"}
            for t in item.get("items", [])
        ]
        return self.todo_block(self._thread_id or "codex", todos)

    async def _parse_event(self, event: dict) -> AsyncIterator[Block]:
        etype = event.get("type", "")

        if etype == "thread.started":
            # codex 0.139 emits `thread_id`; older builds nested `thread:{id}`.
            tid = event.get("thread_id") or (event.get("thread") or {}).get("id")
            if tid and tid != self._thread_id:
                self._thread_id = tid
                if self._on_session_id:
                    self._on_session_id(tid)
            return

        if etype == "item.started":
            item = event.get("item", {})
            self._pending_items[item.get("id", "")] = item
            itype = item.get("type", "")
            if itype in ("command_execution", "mcp_tool_call", "file_change", "web_search"):
                yield AudioHintBlock(state="agent_silent")
                yield ToolCallBlock(
                    id=item.get("id") or str(uuid.uuid4()),
                    name=_item_tool_name(item),
                    args=_item_tool_args(item),
                    status="running",
                )
            elif itype == "todo_list":
                # codex `update_plan` → a live plan card (upserted in place).
                yield self._todo_card(item)

        elif etype == "item.updated":
            # codex re-sends the whole `todo_list` as a step's status changes.
            item = event.get("item", {})
            if item.get("type") == "todo_list":
                yield self._todo_card(item)

        elif etype == "item.delta":
            item = event.get("item", {})
            itype = item.get("type", "")
            if itype == "agent_message":
                delta = item.get("delta", "")
                if delta:
                    yield TextBlock(text=delta, done=False)

        elif etype == "item.completed":
            item = event.get("item", {})
            item_id = item.get("id", "")
            itype = item.get("type", "")
            self._pending_items.pop(item_id, None)

            if itype == "agent_message":
                # codex 0.139 carries the full text in `item.text` and (in exec --json)
                # emits no item.delta, so emit the full text here. `content` is the older
                # field name. (If deltas ever stream, iOS upserts by the done flag.)
                text = item.get("text") or item.get("content") or ""
                if text:
                    yield TextBlock(text=text, done=True)

            elif itype == "reasoning":
                summary = item.get("summary") or item.get("content", "")
                if summary:
                    yield ThinkingBlock(thinking=summary)

            elif itype == "todo_list":
                yield self._todo_card(item)

            elif itype in ("command_execution", "mcp_tool_call", "file_change", "web_search"):
                is_error = item.get("status") == "error"
                tc_id = item_id or str(uuid.uuid4())
                yield ToolCallBlock(
                    id=tc_id,
                    name=_item_tool_name(item),
                    args=_item_tool_args(item),
                    status="error" if is_error else "done",
                )
                raw_output = item.get("output") or item.get("result")
                if raw_output is not None:
                    content = raw_output if isinstance(raw_output, str) else json.dumps(raw_output)
                    yield tool_result_or_large(
                        tool_call_id=tc_id,
                        content=content,
                        fetch_store=self._fetch_store,
                        exit_code=item.get("exit_code"),
                        is_error=is_error,
                    )
                yield AudioHintBlock(state="agent_speaking")

        elif etype == "turn.completed":
            usage = event.get("usage") or {}
            if usage:
                self._usage = {
                    "input_tokens": usage.get("input_tokens", usage.get("inputTokens", 0)),
                    "output_tokens": usage.get("output_tokens", usage.get("outputTokens", 0)),
                    "cache_read_tokens": usage.get(
                        "cache_read_tokens",
                        usage.get("cacheReadTokens", usage.get("cache_read_input_tokens", 0)),
                    ),
                    "cache_write_tokens": usage.get(
                        "cache_write_tokens",
                        usage.get("cacheWriteTokens", usage.get("cache_creation_input_tokens", 0)),
                    ),
                }
                yield SessionStatusHintBlock(
                    model=self._model or "",
                    output_tokens=self._usage.get("output_tokens", 0),
                    last_input_tokens=self._usage.get("input_tokens", 0),
                    cache_creation_tokens=self._usage.get("cache_write_tokens", 0),
                )
            yield AudioHintBlock(state="agent_silent")
            yield SessionEndBlock(reason="completed")

        elif etype == "turn.failed":
            err = event.get("error", {})
            msg = err.get("message", "Codex turn failed") if isinstance(err, dict) else str(err)
            yield ErrorBlock(message=msg, recoverable=True)
            yield AudioHintBlock(state="agent_silent")
            yield SessionEndBlock(reason="error")


# ── Helpers ───────────────────────────────────────────────────────────────────

def _item_tool_name(item: dict) -> str:
    itype = item.get("type", "")
    if itype == "command_execution":
        return "shell"
    if itype == "mcp_tool_call":
        return item.get("tool", "mcp_tool")
    if itype == "file_change":
        return f"file_{item.get('change_type', 'change')}"
    if itype == "web_search":
        return "web_search"
    return itype


def _item_tool_args(item: dict) -> dict:
    itype = item.get("type", "")
    if itype == "command_execution":
        return {"command": item.get("command", "")}
    if itype == "mcp_tool_call":
        return item.get("params", item.get("arguments", {}))
    if itype == "file_change":
        return {"path": item.get("path", ""), "type": item.get("change_type", "")}
    if itype == "web_search":
        return {"query": item.get("query", "")}
    return {}
