"""Benchmark runner: executes all strategies across all tasks."""

from __future__ import annotations

import json
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path

from cohort.benchmark.config import CodingTask, TaskMetrics
from cohort.benchmark.llm_backend import LLMBackend
from cohort.benchmark.metrics import compute_metrics
from cohort.benchmark.strategies import ALL_STRATEGIES, Strategy
from cohort.benchmark.task_loader import load_all_tasks

logger = logging.getLogger(__name__)


@dataclass
class TaskResult:
    """Results for one task across all strategies."""

    task_id: str
    title: str = ""
    metrics: dict[str, TaskMetrics] = field(default_factory=dict)

    def to_dict(self) -> dict:
        return {
            "task_id": self.task_id,
            "title": self.title,
            "configurations": {
                name: m.to_dict() for name, m in self.metrics.items()
            },
        }


@dataclass
class BenchmarkReport:
    """Full benchmark report across all tasks and strategies."""

    benchmark_id: str
    model: str
    timestamp: str
    task_results: list[TaskResult] = field(default_factory=list)
    total_elapsed_seconds: float = 0.0

    def aggregate(self) -> dict[str, dict]:
        """Compute aggregate metrics per strategy."""
        agg: dict[str, dict] = {}
        for tr in self.task_results:
            for strategy_name, m in tr.metrics.items():
                if strategy_name not in agg:
                    agg[strategy_name] = {
                        "tasks": 0,
                        "total_correct": 0.0,
                        "total_tokens": 0,
                        "total_api_calls": 0,
                        "total_turns": 0,
                        "total_elapsed": 0.0,
                        "total_loops": 0,
                        "total_duplications": 0,
                        "planning_paralysis_count": 0,
                        "fully_correct": 0,
                    }
                a = agg[strategy_name]
                a["tasks"] += 1
                a["total_correct"] += m.correctness
                a["total_tokens"] += m.total_tokens
                a["total_api_calls"] += m.api_calls
                a["total_turns"] += m.turn_count
                a["total_elapsed"] += m.elapsed_seconds
                a["total_loops"] += m.loop_count
                a["total_duplications"] += m.duplication_count
                if m.planning_paralysis:
                    a["planning_paralysis_count"] += 1
                if m.correctness == 1.0:
                    a["fully_correct"] += 1

        # Compute averages
        for name, a in agg.items():
            n = a["tasks"]
            a["avg_correctness"] = round(a["total_correct"] / n, 4) if n else 0
            a["avg_tokens"] = round(a["total_tokens"] / n) if n else 0
            a["avg_turns"] = round(a["total_turns"] / n, 1) if n else 0
            a["success_rate"] = round(a["fully_correct"] / n, 4) if n else 0
            a["failure_rate"] = round(1 - a["success_rate"], 4)

        return agg

    def to_dict(self) -> dict:
        return {
            "benchmark_id": self.benchmark_id,
            "model": self.model,
            "timestamp": self.timestamp,
            "total_elapsed_seconds": round(self.total_elapsed_seconds, 2),
            "tasks": [tr.to_dict() for tr in self.task_results],
            "aggregate": self.aggregate(),
        }

    def to_summary_table(self) -> str:
        """Format results as a readable ASCII table."""
        agg = self.aggregate()
        lines = [
            f"\n  Benchmark: {self.benchmark_id}",
            f"  Model: {self.model}",
            f"  Tasks: {len(self.task_results)}",
            f"  Total time: {self.total_elapsed_seconds:.0f}s",
            "",
            "  Strategy               Success  Avg Correct"
            "  Avg Tokens  Avg Turns  Loops  Plan-Paralysis",
            "  " + "-" * 95,
        ]
        for name in ["single_agent", "round_robin", "hierarchical", "cohort_scored"]:
            if name not in agg:
                continue
            a = agg[name]
            lines.append(
                f"  {name:24s}"
                f"  {a['success_rate']*100:5.1f}%"
                f"      {a['avg_correctness']*100:5.1f}%"
                f"       {a['avg_tokens']:6d}"
                f"       {a['avg_turns']:4.1f}"
                f"      {a['total_loops']:2d}"
                f"             {a['planning_paralysis_count']:2d}"
            )
        lines.append("")
        return "\n".join(lines)


class BenchmarkRunner:
    """Runs all tasks through all strategies and collects results."""

    def __init__(
        self,
        llm: LLMBackend,
        tasks: list[CodingTask],
        output_dir: Path | None = None,
        strategies: list[type[Strategy]] | None = None,
        distill_llm: LLMBackend | None = None,
    ) -> None:
        self.llm = llm
        self.distill_llm = distill_llm
        self.tasks = tasks
        self.output_dir = output_dir or Path("benchmark_results")
        self.strategies = strategies or ALL_STRATEGIES

    def _run_one(
        self, task: CodingTask, strategy_cls: type[Strategy],
    ) -> tuple[str, str, TaskMetrics]:
        """Run a single (task, strategy) pair. Thread-safe."""
        strategy_name = strategy_cls.NAME
        try:
            strategy = strategy_cls(llm=self.llm, task=task)
            # Wire separate distill backend for CohortScored
            if hasattr(strategy, "distill_backend") and self.distill_llm:
                strategy.distill_backend = self.distill_llm
            result = strategy.run()
            metrics = compute_metrics(result, task)
            logger.info(
                "  %s × %s: %.0f%% correct, %d turns, %d tokens",
                task.task_id, strategy_name,
                metrics.correctness * 100,
                metrics.turn_count,
                metrics.total_tokens,
            )
            return (task.task_id, strategy_name, metrics)
        except Exception as e:
            logger.error("  %s × %s FAILED: %s", task.task_id, strategy_name, e)
            return (task.task_id, strategy_name, TaskMetrics(
                strategy_name=strategy_name, task_id=task.task_id,
            ))

    def run_all(self, max_workers: int = 2) -> BenchmarkReport:
        """Run every task through all configured strategies.

        Uses thread parallelism to saturate multiple GPUs. Ollama
        schedules concurrent requests across available devices.
        """
        from concurrent.futures import ThreadPoolExecutor, as_completed

        bench_id = f"bench_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        report = BenchmarkReport(
            benchmark_id=bench_id,
            model=self.llm.model_name,
            timestamp=datetime.now().isoformat(),
        )

        t0 = time.time()

        # Build all (task, strategy) pairs
        jobs = [
            (task, strategy_cls)
            for task in self.tasks
            for strategy_cls in self.strategies
        ]
        total = len(jobs)
        logger.info("Dispatching %d jobs across %d workers", total, max_workers)

        # Collect results keyed by task_id
        results_by_task: dict[str, TaskResult] = {}
        for task in self.tasks:
            results_by_task[task.task_id] = TaskResult(
                task_id=task.task_id, title=task.title,
            )

        completed = 0
        with ThreadPoolExecutor(max_workers=max_workers) as pool:
            futures = {
                pool.submit(self._run_one, task, sc): (task.task_id, sc.NAME)
                for task, sc in jobs
            }
            for future in as_completed(futures):
                task_id, strategy_name, metrics = future.result()
                results_by_task[task_id].metrics[strategy_name] = metrics
                completed += 1
                logger.info("  [%d/%d] done", completed, total)

        report.task_results = list(results_by_task.values())

        report.total_elapsed_seconds = time.time() - t0

        # Save results
        self.output_dir.mkdir(parents=True, exist_ok=True)
        output_path = self.output_dir / f"{bench_id}.json"
        output_path.write_text(
            json.dumps(report.to_dict(), indent=2),
            encoding="utf-8",
        )
        logger.info("Results saved to %s", output_path)

        return report

    @classmethod
    def from_directory(
        cls,
        tasks_dir: Path,
        llm: LLMBackend,
        output_dir: Path | None = None,
    ) -> BenchmarkRunner:
        """Create a runner from a directory of task JSON files."""
        tasks = load_all_tasks(tasks_dir)
        return cls(llm=llm, tasks=tasks, output_dir=output_dir)
