import json
from dataclasses import replace
from pathlib import Path
from uuid import uuid4

from dual_agent.infra.process import ProcessRunner, collect_secrets, run_with_secrets

from ..base import AgentRequest, AgentResult, Reply
from ..capability import Capability, probe_binary
from ..fake import SCHEMAS
from ..parse import classify_failure, find_session_id, parse_codex_stream, read_usage
from ..prompt import stage_prompt

#: Tokens the help output must contain for the adapter to be usable at all.
REQUIRED = {"json": "--json", "schema": "--output-schema", "last_message": "--output-last-message"}

#: Optional capabilities. Missing ones degrade behaviour but never fail the probe.
OPTIONAL = {"resume": "resume"}


def _require_all_properties(schema: object) -> None:
    """Providers accept a subset unless every property is explicitly required."""
    if isinstance(schema, dict):
        if isinstance(schema.get("properties"), dict):
            schema["required"] = list(schema["properties"])
        for value in schema.values():
            _require_all_properties(value)
    elif isinstance(schema, list):
        for value in schema:
            _require_all_properties(value)


class CodexAdapter:
    def __init__(
        self,
        binary: str = "codex",
        log_root: str | Path = ".dual-agent/logs",
        timeout_s: float = 900,
        model: str = "gpt-5.6-luna",
        effort: str = "",
        review_sandbox: str = "read-only",
        allow_network: bool = True,
        passthrough_env: tuple[str, ...] = (),
    ) -> None:
        self.binary = binary
        self.log_root = Path(log_root)
        self.timeout_s = timeout_s
        self.model = model
        #: Passed as `-c model_reasoning_effort=<level>`; empty leaves the CLI default.
        self.effort = effort
        self.review_sandbox = review_sandbox
        self.allow_network = allow_network
        self.passthrough_env = tuple(passthrough_env)
        self._capability: Capability | None = None

    def probe(self) -> Capability:
        self._capability = probe_binary(
            [self.binary], [["--version"], ["exec", "--help"]], {**REQUIRED, **OPTIONAL}, required=set(REQUIRED)
        )
        return self._capability

    def _tuning(self) -> list[str]:
        """Model and reasoning effort, in the form this CLI accepts."""
        argv = ["--model", self.model] if self.model else []
        if self.effort:
            argv += ["-c", f"model_reasoning_effort={self.effort}"]
        return argv

    def supports_resume(self) -> bool:
        capability = self._capability or self.probe()
        return capability.features.get("resume", False)

    def ask(self, prompt: str, cwd: Path, session_id: str | None = None, on_start=None) -> Reply:
        """Answer a question in prose. No schema, no artifact, no state change."""
        self.log_root.mkdir(parents=True, exist_ok=True)
        key = uuid4().hex
        output_path = self.log_root / f"ask-{key}.message.txt"
        argv = [
            self.binary, "exec", *self._tuning(), "--json", "--skip-git-repo-check",
            "--sandbox", self.review_sandbox, "-C", str(cwd),
            "--output-last-message", str(output_path),
        ]
        if session_id and self.supports_resume():
            argv += ["resume", session_id]
        argv.append(prompt)
        result = run_with_secrets(ProcessRunner(self.log_root, allow_network=self.allow_network, extra_env=self.passthrough_env),
            argv, cwd, self.timeout_s, name=f"codex-ask-{key}", on_start=on_start, stdout_suffix="jsonl",
            secrets=collect_secrets(self.passthrough_env),
        )
        if result.status == "TIMEOUT":
            return Reply("", ok=False, error="Codex timed out")
        stream = _read(result.stdout_path)
        text = _read(output_path).strip() or _last_message(stream)
        if not text:
            return Reply("", ok=False, error=(_read(result.stderr_path).strip() or "Codex said nothing"))
        return Reply(text, session_id=find_session_id_in(stream), usage=usage_in(stream))

    def run(self, request: AgentRequest, on_start=None) -> AgentResult:
        self.log_root.mkdir(parents=True, exist_ok=True)
        schema_type = SCHEMAS.get(request.stage)
        if schema_type is None:
            return AgentResult("PROTOCOL", error=f"unsupported stage {request.stage}")
        schema_path = self.log_root / f"{request.run_id}.schema.json"
        output_path = self.log_root / f"{request.run_id}.artifact.json"
        schema = schema_type.model_json_schema()
        _require_all_properties(schema)
        schema_path.write_text(json.dumps(schema), encoding="utf-8")
        writes = request.stage in {"IMPLEMENT", "FIX"}
        # `codex exec [OPTIONS] <COMMAND>`: the resume subcommand comes after the options,
        # and it rejects any option repeated after it.
        argv = [
            self.binary,
            "exec",
            *self._tuning(),
            "--json",
            "--skip-git-repo-check",
            "--sandbox",
            "workspace-write" if writes else self.review_sandbox,
            "-C",
            str(request.worktree_path),
            "--output-schema",
            str(schema_path),
            "--output-last-message",
            str(output_path),
        ]
        if request.session_id and self.supports_resume():
            argv += ["resume", request.session_id]
        argv.append(stage_prompt(request))
        result = run_with_secrets(ProcessRunner(self.log_root, allow_network=self.allow_network, extra_env=self.passthrough_env),
            argv, request.worktree_path, self.timeout_s, name=f"codex-{request.run_id}", on_start=on_start,
            stdout_suffix="jsonl", secrets=collect_secrets(self.passthrough_env),
        )
        evidence = {
            "exit_code": result.exit_code,
            "stdout_path": result.stdout_path,
            "stderr_path": result.stderr_path,
            "duration_ms": result.duration_ms,
        }
        if result.status == "TIMEOUT":
            return AgentResult("TIMEOUT", error="Codex timed out", **evidence)
        stream = _read(result.stdout_path)
        artifact_text = _read(output_path)
        parsed = parse_codex_stream(stream, schema_type, artifact_text)
        if parsed.status == "PROTOCOL":
            diagnostics = f"{stream}\n{_read(result.stderr_path)}"
            category = classify_failure(diagnostics)
            if category != "AGENT_FAILURE":
                return AgentResult(category, error=diagnostics.strip(), session_id=parsed.session_id, **evidence)
        if result.status != "SUCCESS" and parsed.status == "SUCCESS":
            category = classify_failure(_read(result.stderr_path)) if hasattr(result, "stderr_path") else "AGENT_FAILURE"
            return AgentResult(
                category, error=f"Codex exited with {result.exit_code}", session_id=parsed.session_id, **evidence
            )
        return replace(parsed, **evidence)


def _read(path: Path) -> str:
    try:
        return path.read_text(encoding="utf-8", errors="replace")
    except OSError:
        return ""


def _events(stream: str):
    for line in stream.splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            yield json.loads(line)
        except json.JSONDecodeError:
            continue


def _last_message(stream: str) -> str:
    """The final prose the agent produced, for the ask path."""
    text = ""
    for event in _events(stream):
        item = event.get("item") if isinstance(event, dict) else None
        for candidate in (event, item):
            if isinstance(candidate, dict):
                value = candidate.get("last_agent_message") or candidate.get("text")
                if isinstance(value, str) and value.strip():
                    text = value
    return text.strip()


def find_session_id_in(stream: str) -> str | None:
    for event in _events(stream):
        found = find_session_id(event)
        if found:
            return found
    return None


def usage_in(stream: str) -> dict | None:
    usage = None
    for event in _events(stream):
        usage = read_usage(event) or usage
    return usage