"""Codex adapter — `codex exec --json`.

**Status: verified against codex-cli 0.145.0 on 2026-07-27.** [Q1](../../docs/12-open-questions.md)
asked whether `codex exec --json` reports token usage or cost at all. It
reports usage; it does not report cost. The observed stream:

    {"type":"thread.started","thread_id":"019fa2b5-..."}
    {"type":"turn.started"}
    {"type":"item.completed","item":{"id":"item_0","type":"error","message":"..."}}
    {"type":"item.completed","item":{"id":"item_1","type":"agent_message","text":"ok"}}
    {"type":"turn.completed","usage":{"input_tokens":16138,"cached_input_tokens":5504,
     "cache_write_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0}}

Three things that run taught us, none of which were guessable:

- **`codex exec` reads stdin.** Left attached it blocks waiting for input, so
  the subprocess is given `stdin=DEVNULL`.
- **It refuses untrusted directories.** Git worktrees are accepted, which is
  what matters here, but a non-repo `cwd` fails outright.
- **An `item.type == "error"` is not necessarily a failure.** A routine
  "skill descriptions were shortened" notice arrives that way on every run.
  Treating error items as fatal would have marked *every* Codex trial failed,
  which is the kind of bug that looks like a config effect.

Because cost is absent, the adapter reports **no cost** rather than zero:

- `cost_usd=None` is never coerced to 0. `--max-spend` cannot budget against an
  agent that does not report cost, and the run says so out loud instead of
  quietly treating every trial as free.
- Zero tokens is treated by the report as "not reported" rather than as a
  measurement, so a Codex arm cannot appear to be spectacularly cheap when it
  is merely unmeasured.
- **No estimated dollar figure is ever printed for Codex** (Q1's explicit
  instruction). Turn count and wall time remain available and are real
  measurements.

The consequence, if Q1's answer turns out to be "no": the primary metric
(ADR-0010) is unavailable for Codex, and cross-vendor comparison degrades to
pass rate — which docs/05-statistics.md says cannot be resolved at an
affordable N. That is a real limitation of the comparison, not of this file,
and the report surfaces it rather than papering over it.

Working directory comes from the subprocess `cwd`, not a `--cd`/`-C` flag.
Q1 could not confirm such a flag exists, and `cwd` is more robust anyway
(ADR-0003).
"""

from __future__ import annotations

import json
import shutil
import subprocess

from ablate.agents.base import PreflightResult, TrialSpec, TrialTelemetry

#: Field spellings to look for when hunting token counts in the event stream.
#: Deliberately generous: the schema is undocumented, and reading a count that
#: is really there beats reporting none because the key was named differently.
_INPUT_KEYS = ("input_tokens", "prompt_tokens", "input")
_OUTPUT_KEYS = ("output_tokens", "completion_tokens", "output")
_CACHE_READ_KEYS = ("cached_input_tokens", "cache_read_input_tokens", "cache_read")
_CACHE_WRITE_KEYS = ("cache_write_input_tokens", "cache_creation_input_tokens", "cache_write")
#: Billed as output, reported separately.
_REASONING_KEYS = ("reasoning_output_tokens",)


class CodexAdapter:
    name = "codex"

    def preflight(self) -> PreflightResult:
        if not shutil.which("codex"):
            return PreflightResult(ok=False, detail="`codex` not found on PATH")
        proc = subprocess.run(["codex", "--version"], capture_output=True, text=True)
        if proc.returncode != 0:
            return PreflightResult(
                ok=False, detail=proc.stderr.strip() or "codex --version failed"
            )
        return PreflightResult(ok=True, version=proc.stdout.strip())

    def build_command(self, spec: TrialSpec) -> list[str]:
        cmd = ["codex", "exec", "--json"]
        if spec.model:
            cmd += ["--model", spec.model]
        # The config under test is delivered as part of the prompt. Codex has no
        # documented equivalent of --append-system-prompt-file, so injecting it
        # any other way would mean writing an AGENTS.md into the worktree and
        # relying on auto-discovery — which is the ambient-config leak that
        # --bare exists to prevent on the Claude Code side.
        cmd.append(_compose_prompt(spec))
        return cmd

    def invoke(self, spec: TrialSpec) -> TrialTelemetry:
        try:
            proc = subprocess.run(
                self.build_command(spec),
                cwd=spec.cwd,
                capture_output=True,
                text=True,
                timeout=spec.timeout_s,
                # codex exec reads stdin; left attached it blocks forever.
                stdin=subprocess.DEVNULL,
            )
        except subprocess.TimeoutExpired:
            return TrialTelemetry(exit_status=124, error=f"timed out after {spec.timeout_s}s")

        telemetry = parse_events(proc.stdout)
        telemetry = TrialTelemetry(**{**telemetry.__dict__, "exit_status": proc.returncode})

        if proc.returncode != 0 and not telemetry.error:
            return TrialTelemetry(
                **{
                    **telemetry.__dict__,
                    "error": (proc.stderr or proc.stdout)[-4000:]
                    or f"codex exited {proc.returncode}",
                }
            )
        return telemetry


def _compose_prompt(spec: TrialSpec) -> str:
    if spec.config_variant_text is None:
        return spec.prompt
    return f"{spec.config_variant_text.strip()}\n\n---\n\n{spec.prompt}"


def _first_int(payload: dict, keys: tuple[str, ...]) -> int:
    for key in keys:
        value = payload.get(key)
        if isinstance(value, int):
            return value
        if isinstance(value, float):
            return int(value)
    return 0


def _find_usage(node) -> dict | None:  # noqa: ANN001 — arbitrary decoded JSON
    """Depth-first search for a usage-shaped object.

    The event schema is undocumented and has moved before, so this looks for
    the *shape* — a mapping carrying at least one recognizable token count —
    rather than a fixed path that a schema change would silently break.
    """
    if isinstance(node, dict):
        if any(
            key in node
            for key in _INPUT_KEYS + _OUTPUT_KEYS + _CACHE_READ_KEYS + _CACHE_WRITE_KEYS + _REASONING_KEYS
        ):
            return node
        for value in node.values():
            found = _find_usage(value)
            if found is not None:
                return found
    elif isinstance(node, list):
        for value in node:
            found = _find_usage(value)
            if found is not None:
                return found
    return None


def parse_events(stdout: str) -> TrialTelemetry:
    """Normalize a `codex exec --json` event stream into `TrialTelemetry`.

    Written against the observed schema (see the module docstring) but tolerant
    of movement: usage is located by shape, and several field spellings are
    accepted, because the stream is undocumented and has changed before.
    """
    turns = 0
    tool_calls = 0
    session_id: str | None = None
    model: str | None = None
    result = ""
    diagnostics: list[str] = []
    usage: dict | None = None
    saw_turn = False

    for line in stdout.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            event = json.loads(line)
        except json.JSONDecodeError:
            # The CLI prints its own chatter ("Reading additional input from
            # stdin...") alongside the JSON. Not a failure.
            continue
        if not isinstance(event, dict):
            continue

        kind = str(event.get("type") or event.get("event") or "")
        session_id = session_id or event.get("thread_id") or event.get("session_id")
        model = model or event.get("model")

        if kind.endswith("turn.completed"):
            turns += 1
            saw_turn = True
            usage = _find_usage(event) or usage
        elif kind.endswith("item.completed"):
            item = event.get("item") or {}
            item_type = str(item.get("type") or "")
            if item_type == "agent_message":
                result = str(item.get("text") or result)
            elif item_type == "error":
                # Collected, not fatal. A benign skills-budget notice arrives
                # as an error item on every run; failing the trial on it would
                # mark every Codex trial failed.
                diagnostics.append(str(item.get("message") or "codex reported an error"))
            elif item_type not in ("reasoning", ""):
                # command_execution, file_change, mcp_tool_call, web_search…
                tool_calls += 1

    usage = usage or {}
    input_tokens = _first_int(usage, _INPUT_KEYS)
    cache_read = _first_int(usage, _CACHE_READ_KEYS)
    cache_write = _first_int(usage, _CACHE_WRITE_KEYS)
    # Reasoning tokens are billed as output and are reported separately, so
    # dropping them would understate the primary metric on every trial.
    output_tokens = _first_int(usage, _OUTPUT_KEYS) + _first_int(usage, _REASONING_KEYS)

    # Failure is judged by the process and by whether a turn ever completed —
    # never by the presence of an error *item*, which is routine.
    error = None
    if not saw_turn:
        error = "; ".join(diagnostics) or "codex produced no completed turn"

    return TrialTelemetry(
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        cache_read=cache_read,
        cache_write=cache_write,
        # Never derived from token counts: Q1 is explicit that no estimated
        # dollar figure is printed for Codex, and a spend ceiling enforced
        # against a guess is worse than none because it looks like it works.
        cost_usd=usage.get("cost_usd") if isinstance(usage.get("cost_usd"), int | float) else None,
        turns=turns,
        tool_calls=tool_calls,
        session_id=session_id,
        model=model,
        raw_result=result,
        error=error,
    )
