import asyncio
import json
import logging
import uuid
import time
import os
from datetime import datetime

from sqlalchemy.orm import Session

from database import SessionLocal
from models import Task, Agent, Run, RunStatus
from core.sandbox import Sandbox
from adapters import get_adapter
from adapters.base import RunContext
from core.metrics import (
    compute_cost,
    parse_pytest_output,
    compute_trajectory_efficiency,
    analyze_trajectory,
    compute_diff_minimality,
    compute_composite_score,
)
from core.judge import Judge, JudgeResult
from core.alignment import ALIGNMENT_RUBRIC_VERSION

logger = logging.getLogger(__name__)


class Runner:
    def __init__(self):
        self.judge = Judge()
        self._active_runs: dict[int, dict] = {}  # run_id -> progress info

    async def run_single(self, db: Session, task: Task, agent: Agent) -> Run:
        """Execute a single benchmark run."""
        agent_config = json.loads(agent.config_json or "{}")
        replicate_index = self._next_replicate_index(db, task.id, agent.id)
        # 1. Create Run record
        run = Run(
            task_id=task.id,
            agent_id=agent.id,
            status=RunStatus.running,
            task_family=self._task_family(task),
            seed=str(agent_config.get("seed", "")) or None,
            replicate_index=replicate_index,
            experiment_id=agent_config.get("experiment_id"),
            arm_id=agent_config.get("arm_id") or agent.name,
            alignment_rubric_version=ALIGNMENT_RUBRIC_VERSION,
            created_at=datetime.utcnow(),
        )
        db.add(run)
        db.commit()
        db.refresh(run)

        self._active_runs[run.id] = {"status": "starting", "started_at": time.time()}
        sandbox = None
        judge_sandbox = None
        start_time = time.time()

        try:
            # 2. Create agent sandbox
            self._active_runs[run.id]["status"] = "creating_sandbox"
            sandbox = await asyncio.to_thread(Sandbox.create, task)

            # 3. Get adapter and build RunContext
            self._active_runs[run.id]["status"] = "running_agent"
            adapter = get_adapter(agent.adapter_type, agent_config)

            ctx = RunContext(
                work_dir="/workspace",
                task_prompt=task.description or task.title,
                success_cmd=task.success_cmd,
                time_limit_s=task.time_limit_s,
                token_limit=task.token_limit,
                repo_url=task.repo_url or "",
                repo_ref=task.repo_ref or "main",
            )

            # 4. Run adapter — returns RunResult dataclass
            adapter_result = await adapter.run(ctx, sandbox)

            # 5. Get diff — prefer adapter-provided diff (e.g. ClaudeCode runs on host),
            #    fall back to sandbox.get_diff() for container-based adapters.
            self._active_runs[run.id]["status"] = "collecting_diff"
            if adapter_result.diff_text:
                diff_text = adapter_result.diff_text
            else:
                diff_text = await asyncio.to_thread(sandbox.get_diff)

            # 5b. Destroy agent container — judge must not share it
            await asyncio.to_thread(sandbox.destroy)
            sandbox = None

            # 6. Create isolated judge container (fresh clone + patch applied)
            self._active_runs[run.id]["status"] = "creating_judge_sandbox"
            judge_sandbox = await asyncio.to_thread(
                Sandbox.create_judge_sandbox, task, diff_text
            )

            # Abort if patch failed to apply — running judge on unpatched repo gives false negatives
            if diff_text and diff_text.strip() and not judge_sandbox.patch_applied:
                patch_output = judge_sandbox.patch_output or "unknown"
                run.invalid_patch_flag = 1
                raise RuntimeError(f"Patch failed to apply in judge container: {patch_output}")

            # 7. Judge in the isolated container
            self._active_runs[run.id]["status"] = "judging"
            judge_result = await asyncio.to_thread(
                self.judge.judge_run, judge_sandbox, task.success_cmd
            )

            # 7. Compute metrics
            # Use adapter-measured duration (only CLI time), not wall-clock
            # which includes sandbox creation + judge overhead.
            agent_duration_s = adapter_result.duration_s
            end_to_end_duration_s = time.time() - start_time
            tokens_input = adapter_result.tokens_input
            tokens_output = adapter_result.tokens_output
            cost_usd = adapter_result.cost_usd

            trajectory = adapter_result.trajectory

            # Trajectory and diff analysis
            traj_metrics = analyze_trajectory(trajectory)
            diff_metrics = compute_diff_minimality(diff_text)

            functional_success = 1 if judge_result.functional_score >= 1.0 else 0
            quality_score = (
                judge_result.regression_score * 0.5
                + judge_result.lint_score * 0.25
                + judge_result.complexity_score * 0.25
            )
            final_status = RunStatus.success if functional_success else RunStatus.failure
            tests_total = judge_result.tests_total if judge_result.tests_total > 0 else adapter_result.tests_total
            tests_passed = judge_result.tests_passed if judge_result.tests_total > 0 else adapter_result.tests_passed
            test_pass_rate = (tests_passed / tests_total) if tests_total > 0 else None

            # Build full metrics JSON
            judge_metrics = {
                "functional_score": judge_result.functional_score,
                "functional_success": functional_success,
                "regression_score": judge_result.regression_score,
                "regression_tests_total": judge_result.regression_tests_total,
                "regression_tests_passed": judge_result.regression_tests_passed,
                "lint_score": judge_result.lint_score,
                "lint_delta": judge_result.lint_delta,
                "complexity_score": judge_result.complexity_score,
                "complexity_delta": judge_result.complexity_delta,
                "quality_score": quality_score,
            }

            # Normalize cost to 0-1 efficiency score (lower cost = higher score)
            efficiency = 1.0 / (1.0 + cost_usd / 0.50) if cost_usd > 0 else 1.0

            composite_score = compute_composite_score({
                "functionality": judge_result.functional_score,
                "directness_score": traj_metrics.get("directness_score", 0),
                "efficiency": efficiency,
                "self_correction_rate": traj_metrics.get("self_correction_rate", 0),
                "churn": diff_metrics.get("churn", 0),
                "loop_count": traj_metrics.get("loop_count", 0),
                "total_steps": traj_metrics.get("total_steps", 1),
            })

            full_metrics = {
                "trajectory": traj_metrics,
                "diff": diff_metrics,
                "judge": judge_metrics,
                "composite_score": composite_score,
                "eval": {
                    "functional_success": functional_success,
                    "test_pass_rate": test_pass_rate,
                    "tests_total": tests_total,
                    "tests_passed": tests_passed,
                    "agent_duration_s": agent_duration_s,
                    "end_to_end_duration_s": round(end_to_end_duration_s, 2),
                    "timeout_flag": 1 if adapter_result.status == RunStatus.timeout else 0,
                    "tool_failure_flag": 1 if adapter_result.errors_count > 0 else 0,
                    "invalid_patch_flag": 0,
                    "task_family": run.task_family,
                    "seed": run.seed,
                    "replicate_index": replicate_index,
                    "experiment_id": run.experiment_id,
                    "arm_id": run.arm_id,
                    "alignment_rubric_version": ALIGNMENT_RUBRIC_VERSION,
                },
            }

            # 8. Update Run record
            run.status = final_status
            run.duration_s = agent_duration_s
            run.agent_duration_s = agent_duration_s
            run.end_to_end_duration_s = round(end_to_end_duration_s, 2)
            run.tokens_input = tokens_input
            run.tokens_output = tokens_output
            run.cost_usd = cost_usd
            run.tests_total = tests_total
            run.tests_passed = tests_passed
            run.functional_success = functional_success
            run.test_pass_rate = test_pass_rate
            run.quality_score = round(quality_score, 4)
            run.errors_count = adapter_result.errors_count
            run.tool_calls = adapter_result.tool_calls
            run.diff_text = diff_text
            run.trajectory = json.dumps(trajectory)
            run.judge_score = judge_result.overall_score
            run.judge_notes = judge_result.overall_notes
            run.timeout_flag = 1 if adapter_result.status == RunStatus.timeout else 0
            run.tool_failure_flag = 1 if adapter_result.errors_count > 0 else 0
            run.invalid_patch_flag = 0
            run.metrics_json = json.dumps(full_metrics)
            db.commit()
            db.refresh(run)

            # 8b. Fire-and-forget LLM judge (non-blocking)
            if os.environ.get("MINIMAX_API_KEY"):
                asyncio.create_task(
                    _run_llm_judges(db, run.id, task, agent, full_metrics, diff_text, adapter_result.raw_output)
                )

        except asyncio.TimeoutError:
            run.status = RunStatus.timeout
            run.duration_s = time.time() - start_time
            run.agent_duration_s = run.duration_s
            run.end_to_end_duration_s = run.duration_s
            run.judge_notes = f"Run timed out after {task.time_limit_s}s"
            run.timeout_flag = 1
            db.commit()
        except Exception as e:
            run.status = RunStatus.error
            run.duration_s = time.time() - start_time
            run.agent_duration_s = run.duration_s
            run.end_to_end_duration_s = run.duration_s
            run.judge_notes = f"Error: {str(e)}"
            run.tool_failure_flag = 1
            db.commit()
        finally:
            # 9. Destroy sandboxes (agent may already be gone, judge always cleaned up)
            if sandbox is not None:
                await asyncio.to_thread(sandbox.destroy)
            if judge_sandbox is not None:
                await asyncio.to_thread(judge_sandbox.destroy)
            self._active_runs.pop(run.id, None)

        db.refresh(run)
        return run

    async def _run_single_for_k(
        self, db: Session, task: Task, agent: Agent, group_id: str
    ) -> Run:
        """Run a single instance for pass@k, tagging with group_id."""
        run = await self.run_single(db, task, agent)
        run.pass_at_k_group = group_id
        db.commit()
        db.refresh(run)
        return run

    async def run_pass_at_k(
        self, db: Session, task: Task, agent: Agent, k: int = 3
    ) -> list[Run]:
        """Run k parallel benchmark runs for pass@k calculation."""
        group_id = str(uuid.uuid4())[:8]
        task_id, agent_id = task.id, agent.id

        async def _single(i):
            run_db = SessionLocal()
            try:
                t = run_db.query(Task).filter(Task.id == task_id).first()
                a = run_db.query(Agent).filter(Agent.id == agent_id).first()
                run = await self.run_single(run_db, t, a)
                run.pass_at_k_group = group_id
                run_db.commit()
                run_db.refresh(run)
                return run
            finally:
                run_db.close()

        results = await asyncio.gather(*[_single(i) for i in range(k)], return_exceptions=True)
        # Log exceptions but don't crash the entire batch
        clean = []
        for i, r in enumerate(results):
            if isinstance(r, Exception):
                logger.error("pass@k run %d/%d failed: %s", i + 1, k, r)
            else:
                clean.append(r)
        return clean

    def get_progress(self, run_id: int) -> dict | None:
        return self._active_runs.get(run_id)

    def _task_family(self, task: Task) -> str:
        slug = task.slug or ""
        if "_" in slug:
            return slug.split("_", 1)[0]
        if "-" in slug:
            return slug.split("-", 1)[0]
        return slug or "unknown"

    def _next_replicate_index(self, db: Session, task_id: int, agent_id: int) -> int:
        return (
            db.query(Run)
            .filter(Run.task_id == task_id, Run.agent_id == agent_id)
            .count()
            + 1
        )


async def _run_llm_judges(
    db: Session, run_id: int, task: Task, agent: Agent,
    full_metrics: dict, diff_text: str, raw_output: str,
) -> None:
    """Fire-and-forget LLM judges — updates metrics_json async."""
    try:
        from core.llm_judge import llm_judge
        from core.alignment import llm_alignment_judge

        judge_metrics = full_metrics.get("judge", {})
        run_data = {
            "task_title": task.title,
            "task_description": task.description or "",
            "agent_name": agent.name,
            "adapter_type": agent.adapter_type,
            "status": "success" if judge_metrics.get("functional_score", 0) >= 1.0 else "failure",
            "functional_score": judge_metrics.get("functional_score"),
            "functional_success": judge_metrics.get("functional_success"),
            "regression_score": judge_metrics.get("regression_score"),
            "lint_score": judge_metrics.get("lint_score"),
            "complexity_score": judge_metrics.get("complexity_score"),
            "composite_score": full_metrics.get("composite_score"),
            "trajectory": full_metrics.get("trajectory", {}),
            "diff_text": diff_text or "",
            "judge_notes": "",
            "tests_total": full_metrics.get("eval", {}).get("tests_total", 0),
            "tests_passed": full_metrics.get("eval", {}).get("tests_passed", 0),
            "raw_output": raw_output or "",
        }

        quality_result = await llm_judge(run_data)
        alignment_result = await llm_alignment_judge(run_data)
        if quality_result:
            full_metrics["llm_judge"] = quality_result
        if alignment_result:
            full_metrics["alignment"] = alignment_result
        if quality_result or alignment_result:
            from database import SessionLocal
            judge_db = SessionLocal()
            try:
                from models import Run
                r = judge_db.query(Run).filter(Run.id == run_id).first()
                if r:
                    import json
                    if alignment_result:
                        r.alignment_score_llm = alignment_result.get("alignment_score")
                        r.alignment_rubric_version = alignment_result.get("rubric_version")
                    r.metrics_json = json.dumps(full_metrics)
                    judge_db.commit()
            finally:
                judge_db.close()
    except Exception as e:
        import logging
        logging.getLogger(__name__).warning("LLM judge background task failed: %s", e)