from __future__ import annotations

import asyncio
import json
from dataclasses import dataclass, field
from pathlib import Path


@dataclass
class CodexResult:
    ok: bool
    final_message: str
    session_id: str | None
    commands: list[str] = field(default_factory=list)
    stderr: str = ""
    returncode: int | None = None


class CodexRunner:
    def __init__(
        self,
        codex_bin: str,
        timeout_seconds: int,
        resume_json_flag_style: str,
    ) -> None:
        self.codex_bin = codex_bin
        self.timeout_seconds = timeout_seconds
        self.resume_json_flag_style = resume_json_flag_style

    def _command(
        self,
        prompt: str,
        sandbox: str,
        session_id: str | None,
        skip_git_repo_check: bool,
    ) -> list[str]:
        base = [self.codex_bin, "exec", "--json", "--sandbox", sandbox]
        if skip_git_repo_check:
            base.append("--skip-git-repo-check")
        if not session_id:
            return [*base, prompt]

        if self.resume_json_flag_style == "after_resume":
            return [
                self.codex_bin,
                "exec",
                "resume",
                session_id,
                "--json",
                "--sandbox",
                sandbox,
                *(["--skip-git-repo-check"] if skip_git_repo_check else []),
                prompt,
            ]
        return [*base, "resume", session_id, prompt]

    async def run(
        self,
        prompt: str,
        cwd: Path,
        sandbox: str,
        session_id: str | None,
        skip_git_repo_check: bool = False,
    ) -> CodexResult:
        cmd = self._command(
            prompt=prompt,
            sandbox=sandbox,
            session_id=session_id,
            skip_git_repo_check=skip_git_repo_check,
        )
        proc = await asyncio.create_subprocess_exec(
            *cmd,
            cwd=str(cwd),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(),
                timeout=self.timeout_seconds,
            )
        except asyncio.TimeoutError:
            proc.kill()
            await proc.wait()
            return CodexResult(
                ok=False,
                final_message=f"Codex timed out after {self.timeout_seconds} seconds.",
                session_id=session_id,
            )

        parsed = self._parse_jsonl(stdout.decode("utf-8", errors="replace"))
        stderr_text = stderr.decode("utf-8", errors="replace").strip()
        if proc.returncode != 0:
            message = parsed.final_message or stderr_text or "Codex command failed."
            return CodexResult(
                ok=False,
                final_message=message,
                session_id=parsed.session_id or session_id,
                commands=parsed.commands,
                stderr=stderr_text,
                returncode=proc.returncode,
            )

        return CodexResult(
            ok=True,
            final_message=parsed.final_message or "(Codex finished without text output.)",
            session_id=parsed.session_id or session_id,
            commands=parsed.commands,
            stderr=stderr_text,
            returncode=proc.returncode,
        )

    @staticmethod
    def _parse_jsonl(text: str) -> CodexResult:
        session_id: str | None = None
        final_messages: list[str] = []
        commands: list[str] = []
        fallback_lines: list[str] = []

        for line in text.splitlines():
            line = line.strip()
            if not line:
                continue
            try:
                event = json.loads(line)
            except json.JSONDecodeError:
                fallback_lines.append(line)
                continue

            if event.get("type") == "thread.started":
                session_id = event.get("thread_id") or session_id
                continue

            item = event.get("item")
            if isinstance(item, dict):
                if item.get("type") == "agent_message" and item.get("text"):
                    final_messages.append(str(item["text"]))
                if item.get("type") == "command_execution" and item.get("command"):
                    commands.append(str(item["command"]))

        final_message = final_messages[-1] if final_messages else "\n".join(fallback_lines)
        return CodexResult(
            ok=True,
            final_message=final_message,
            session_id=session_id,
            commands=commands,
        )