import json
import logging
import os
import re
import shlex
import shutil
import subprocess
import tempfile
import time

from config import CODEX_CLI_PATH, get_pricing
from models import RunStatus

from adapters.base import AgentAdapter, RunContext, RunResult

logger = logging.getLogger(__name__)

# Codex prints token usage like: "tokens used\n1.019"
_TOKENS_USED_RE = re.compile(r"tokens\s+used\s*\n\s*([0-9.]+)", re.IGNORECASE)
# Broader alternatives
_TOKEN_ALT_RE = re.compile(r"(?:total\s+)?tokens[:\s]+(\d[\d,]*)", re.I)
_COST_ALT_RE = re.compile(r"(?:cost|price|total)[:\s]*\$([0-9.]+)", re.I)

# Default model assumed when Codex doesn't report one
_DEFAULT_MODEL = "gpt-4o"


class CodexAdapter(AgentAdapter):
    def name(self) -> str:
        return "Codex"

    def is_available(self) -> bool:
        return shutil.which(CODEX_CLI_PATH) is not None

    async def run(self, ctx: RunContext, sandbox) -> RunResult:
        """Run Codex CLI on the host against a local clone of the task repo.

        Codex uses ChatGPT Plus OAuth (no API key needed).
        Runs non-interactively via `codex exec -s full-auto`.
        """
        start = time.time()
        tmpdir = tempfile.mkdtemp(prefix="arena-codex-")

        try:
            # Clone repo to host temp dir.
            repo_url = ctx.repo_url or ""
            repo_ref = ctx.repo_ref or "main"

            if repo_url:
                clone_result = subprocess.run(
                    ["git", "clone", "--depth=1", "-b", repo_ref, repo_url, tmpdir],
                    capture_output=True, text=True, timeout=120,
                )
                if clone_result.returncode != 0:
                    return RunResult(
                        status=RunStatus.error,
                        duration_s=round(time.time() - start, 2),
                        raw_output=clone_result.stderr[:5000],
                    )
                work_dir = tmpdir
            else:
                work_dir = ctx.work_dir

            # Git checkpoint for diff.
            subprocess.run(
                ["git", "tag", "arena-checkpoint"],
                cwd=work_dir, capture_output=True, timeout=10,
            )

            # Use model from agent config, fall back to default.
            model = self.config.get("model", _DEFAULT_MODEL)

            prompt = (
                f"Fix the following issue in the codebase:\n\n"
                f"{ctx.task_prompt}\n\n"
                f"After fixing, verify your changes pass: {ctx.success_cmd}"
            )

            cmd = [
                CODEX_CLI_PATH,
                "exec",
                "-s", "danger-full-access",
                "--ephemeral",
                "--model", model,
                prompt,
            ]

            try:
                proc = subprocess.run(
                    cmd,
                    cwd=work_dir,
                    capture_output=True,
                    text=True,
                    timeout=ctx.time_limit_s,
                )
                returncode = proc.returncode
                stdout = proc.stdout
                stderr = proc.stderr
            except subprocess.TimeoutExpired:
                returncode = 124
                stdout = ""
                stderr = f"Command timed out after {ctx.time_limit_s}s"

            duration = time.time() - start

            # Collect diff, exclude __pycache__ and .pyc.
            def _get_diff(rev):
                r = subprocess.run(
                    ["git", "diff", rev, "--", ".", ":!**/__pycache__/", ":!**/*.pyc"],
                    cwd=work_dir, capture_output=True, text=True, timeout=30,
                )
                return r.stdout if r.returncode == 0 else ""

            diff_text = _get_diff("arena-checkpoint")
            if not diff_text.strip():
                diff_text = _get_diff("arena-checkpoint..HEAD")

            # Parse token usage from output.
            tokens_input = 0
            tokens_output = 0
            cost_usd = 0.0
            combined = stdout + "\n" + stderr

            # Log raw output for diagnostics
            logger.info("CODEX RAW OUTPUT (first 2000 chars):\n%s", combined[:2000])

            # Codex prints "tokens used\n1.019" (in thousands)
            m = _TOKENS_USED_RE.search(combined)
            if m:
                total_k = float(m.group(1))
                total_tokens = int(total_k * 1000)
                # Rough split: 70% input, 30% output
                tokens_input = int(total_tokens * 0.7)
                tokens_output = total_tokens - tokens_input

            # Fallback: try alternative token/cost patterns
            if tokens_input == 0 and tokens_output == 0:
                # Filter out metadata lines (e.g., "Repo-map: using 1024 tokens")
                content_lines = [
                    line for line in combined.split("\n")
                    if not re.search(r"repo.?map|using \d+ tokens.*refresh", line, re.I)
                ]
                filtered = "\n".join(content_lines)
                alt_tokens = _TOKEN_ALT_RE.findall(filtered)
                if alt_tokens:
                    total = sum(int(t.replace(",", "")) for t in alt_tokens)
                    if total >= 100:  # Sanity check: real LLM usage > 100 tokens
                        tokens_input = int(total * 0.7)
                        tokens_output = total - tokens_input
            if cost_usd == 0.0:
                alt_costs = _COST_ALT_RE.findall(combined)
                if alt_costs:
                    cost_usd = sum(float(c) for c in alt_costs)

            # Estimate cost from tokens using dynamic pricing lookup.
            if cost_usd == 0.0 and (tokens_input or tokens_output):
                rate_in, _, _, rate_out = get_pricing(model)
                cost_usd = (tokens_input * rate_in + tokens_output * rate_out) / 1_000_000

            # Last resort: estimate from raw output length (better than $0)
            if cost_usd == 0.0 and (tokens_input + tokens_output) < 50 and len(combined) > 100:
                estimated_output_tokens = len(combined) // 4
                rate_in, _, _, rate_out = get_pricing(model)
                tokens_input = estimated_output_tokens * 3
                tokens_output = estimated_output_tokens
                cost_usd = (tokens_input * rate_in + tokens_output * rate_out) / 1_000_000
                logger.warning("Cost estimated from output length (%d chars) — no token data found", len(combined))

            trajectory = [{"tool_name": "Bash", "input_preview": f"codex exec --model {model}"}]

            if duration >= ctx.time_limit_s:
                status = RunStatus.timeout
            elif returncode == 0:
                status = RunStatus.success
            else:
                status = RunStatus.failure

            return RunResult(
                status=status,
                duration_s=round(duration, 2),
                tokens_input=tokens_input,
                tokens_output=tokens_output,
                cost_usd=round(cost_usd, 6),
                diff_text=diff_text,
                trajectory=trajectory,
                raw_output=(stdout + stderr)[:50_000],
            )
        finally:
            shutil.rmtree(tmpdir, ignore_errors=True)