"""Simple evaluation runner for running strategies against benchmarks.

This module provides the core evaluation loop that runs a strategy against
a benchmark and computes metrics. Supports optional checkpointing for
resumption after interruption.
"""

import asyncio
import json
import logging
from pathlib import Path
from typing import Any

from embergraphrouter.evals.benchmarks.base import Benchmark, BenchmarkExample
from embergraphrouter.evals.metrics.aggregation import (
    aggregate_by_slice,
    compute_efficiency_metrics,
    compute_quality_metrics,
    compute_routing_metrics,
)
from embergraphrouter.evals.runners.base import (
    EvaluationResult,
    ExampleResult,
    RunConfig,
    Runner,
)
from embergraphrouter.evals.runners.checkpoint import (
    AdaptiveRateLimiter,
    CheckpointManager,
    ProgressTracker,
    compute_config_hash,
)
from embergraphrouter.evals.strategies.base import Strategy

logger = logging.getLogger(__name__)


class SimpleRunner(Runner):
    """Simple evaluation runner with optional checkpointing.

    Runs a strategy against a benchmark, evaluates each example,
    and computes aggregate metrics. Supports checkpointing for
    resumption after interruption.

    Example:
        # Basic usage
        runner = SimpleRunner(strategy, benchmark)
        result = await runner.run()

        # With checkpointing
        runner = SimpleRunner(
            strategy, benchmark,
            checkpoint_dir="checkpoints",
            parallelism=10,
        )
        result = await runner.run()

        # Resume interrupted run
        runner = SimpleRunner.resume("checkpoints/run_abc123.pkl", strategy, benchmark)
        result = await runner.run()
    """

    def __init__(
        self,
        strategy: Strategy,
        benchmark: Benchmark,
        max_examples: int | None = None,
        system_prompt: str | None = None,
        subset: str | None = None,
        checkpoint_dir: str | Path | None = None,
        parallelism: int = 1,
        checkpoint_interval: int = 50,
        timeout_per_example: float = 120.0,
        checkpoint_manager: CheckpointManager | None = None,
    ):
        """Initialize runner.

        Args:
            strategy: Strategy to evaluate.
            benchmark: Benchmark to run against.
            max_examples: Maximum examples to evaluate (None = all).
            system_prompt: Optional system prompt for all queries.
            subset: Benchmark subset to use.
            checkpoint_dir: Directory for checkpoints (None = no checkpointing).
            parallelism: Number of concurrent evaluations (1 = sequential).
            checkpoint_interval: Save checkpoint every N examples.
            timeout_per_example: Timeout per example in seconds.
            checkpoint_manager: Pre-configured checkpoint manager for resumption.
        """
        self._strategy = strategy
        self._benchmark = benchmark
        self._max_examples = max_examples
        self._system_prompt = system_prompt
        self._subset = subset
        self._checkpoint_dir = Path(checkpoint_dir) if checkpoint_dir else None
        self._parallelism = parallelism
        self._checkpoint_interval = checkpoint_interval
        self._timeout = timeout_per_example
        self._checkpoint_manager = checkpoint_manager
        self._rate_limiter: AdaptiveRateLimiter | None = None

    @classmethod
    def resume(
        cls,
        checkpoint_path: str | Path,
        strategy: Strategy,
        benchmark: Benchmark,
        **kwargs: Any,
    ) -> "SimpleRunner":
        """Resume runner from checkpoint.

        Args:
            checkpoint_path: Path to checkpoint file.
            strategy: Strategy instance.
            benchmark: Benchmark instance.
            **kwargs: Additional runner arguments.

        Returns:
            SimpleRunner configured to resume.
        """
        manager = CheckpointManager.resume(checkpoint_path)

        return cls(
            strategy=strategy,
            benchmark=benchmark,
            checkpoint_manager=manager,
            **kwargs,
        )

    def _get_config_hash(self) -> str:
        """Compute hash of current configuration."""
        return compute_config_hash({
            "strategy": self._strategy.name,
            "benchmark": self._benchmark.name,
            "max_examples": self._max_examples,
            "subset": self._subset,
        })

    def setup(self) -> None:
        """Pre-evaluation setup."""
        self._strategy.warm_up()

        # Initialize checkpointing if enabled
        if self._checkpoint_dir and not self._checkpoint_manager:
            self._checkpoint_manager = CheckpointManager(
                checkpoint_dir=self._checkpoint_dir,
                config_hash=self._get_config_hash(),
                save_interval=self._checkpoint_interval,
            )

        # Initialize rate limiter for parallel execution
        if self._parallelism > 1:
            self._rate_limiter = AdaptiveRateLimiter(
                initial=self._parallelism,
                min_level=1,
                max_level=self._parallelism * 2,
            )

    async def run(self) -> EvaluationResult:
        """Execute the evaluation.

        Supports both sequential and parallel execution, with optional
        checkpointing for resumption.

        Returns:
            EvaluationResult with all metrics and per-example results.
        """
        # Create result container
        result = EvaluationResult.create(
            strategy_name=self._strategy.name,
            benchmark_name=self._benchmark.name,
        )
        result.config = RunConfig().to_dict()

        # Load examples
        examples = list(self._benchmark.load(self._subset))
        if self._max_examples:
            examples = examples[: self._max_examples]

        result.n_examples = len(examples)

        # Check for resumed results
        resumed_count = 0
        if self._checkpoint_manager:
            resumed_count = self._checkpoint_manager.n_completed
            if resumed_count > 0:
                logger.info(
                    f"Resuming from checkpoint: {resumed_count} examples already done"
                )

        logger.info(
            "Running %s on %d examples from %s (parallelism=%d)",
            self._strategy.name,
            result.n_examples,
            self._benchmark.name,
            self._parallelism,
        )

        # Run inference and evaluation
        if self._parallelism > 1 and self._rate_limiter:
            example_results, errors = await self._run_parallel(examples)
        else:
            example_results, errors = await self._run_sequential(examples)

        result.example_results = example_results
        result.errors = errors

        # Compute aggregate metrics
        result.quality = compute_quality_metrics(example_results)
        result.efficiency = compute_efficiency_metrics(example_results)
        result.routing = compute_routing_metrics(example_results)

        # Compute sliced metrics
        result.by_difficulty = aggregate_by_slice(example_results, "difficulty")
        result.by_domain = aggregate_by_slice(example_results, "domain")
        result.by_tier = aggregate_by_slice(example_results, "tier")

        # Final checkpoint save
        if self._checkpoint_manager:
            self._checkpoint_manager.save()

        acc = 0.0
        if result.quality:
            acc_metric = result.quality.get("accuracy") or result.quality.get("score")
            if acc_metric:
                acc = acc_metric.mean

        logger.info(
            "Evaluation complete: accuracy=%.3f, n=%d",
            acc,
            result.n_examples,
        )

        return result

    async def _run_sequential(
        self,
        examples: list[BenchmarkExample],
    ) -> tuple[list[ExampleResult], list[str]]:
        """Run evaluation sequentially."""
        example_results: list[ExampleResult] = []
        errors: list[str] = []
        progress = ProgressTracker(total=len(examples))

        for i, example in enumerate(examples):
            # Skip if already completed (resumption)
            if self._checkpoint_manager and self._checkpoint_manager.is_completed(
                example.id
            ):
                cached = self._checkpoint_manager.get_result(example.id)
                if cached:
                    example_results.append(ExampleResult(**cached))
                progress.increment()
                continue

            try:
                example_result = await self._evaluate_single(example)
                example_results.append(example_result)

                # Checkpoint
                if self._checkpoint_manager:
                    self._checkpoint_manager.mark_completed(
                        example.id,
                        example_result.__dict__,
                    )

                progress.increment()

                if (i + 1) % 10 == 0:
                    logger.info(progress.format_status())

            except Exception as e:
                error_msg = f"Error on example {example.id}: {e}"
                logger.warning(error_msg)
                errors.append(error_msg)
                if self._checkpoint_manager:
                    self._checkpoint_manager.record_error(error_msg)

                example_results.append(
                    ExampleResult(
                        example_id=example.id,
                        strategy_name=self._strategy.name,
                        prediction="",
                        reference=example.reference,
                        correct=False,
                        score=0.0,
                        model_used="error",
                        error=str(e),
                        metadata=example.metadata,
                    )
                )
                progress.increment()

        return example_results, errors

    async def _run_parallel(
        self,
        examples: list[BenchmarkExample],
    ) -> tuple[list[ExampleResult], list[str]]:
        """Run evaluation with adaptive parallelism."""
        assert self._rate_limiter is not None

        example_results: list[ExampleResult] = []
        errors: list[str] = []
        results_lock = asyncio.Lock()
        progress = ProgressTracker(total=len(examples))

        async def process_one(example: BenchmarkExample) -> None:
            # Skip if already completed
            if self._checkpoint_manager and self._checkpoint_manager.is_completed(
                example.id
            ):
                cached = self._checkpoint_manager.get_result(example.id)
                if cached:
                    async with results_lock:
                        example_results.append(ExampleResult(**cached))
                progress.increment()
                return

            ctx = await self._rate_limiter.acquire()
            try:
                async with ctx:
                    try:
                        result = await asyncio.wait_for(
                            self._evaluate_single(example),
                            timeout=self._timeout,
                        )
                        await self._rate_limiter.record_success()
                    except asyncio.TimeoutError:
                        await self._rate_limiter.record_error()
                        result = ExampleResult(
                            example_id=example.id,
                            strategy_name=self._strategy.name,
                            prediction="",
                            reference=example.reference,
                            correct=False,
                            score=0.0,
                            model_used="timeout",
                            error="Timeout",
                            metadata=example.metadata,
                        )

                    async with results_lock:
                        example_results.append(result)
                        if self._checkpoint_manager:
                            self._checkpoint_manager.mark_completed(
                                example.id,
                                result.__dict__,
                            )

                    progress.increment()

            except Exception as e:
                # Check for rate limit errors
                if "rate" in str(e).lower():
                    await self._rate_limiter.record_rate_limit()
                else:
                    await self._rate_limiter.record_error()

                error_msg = f"Error on example {example.id}: {e}"
                async with results_lock:
                    errors.append(error_msg)
                    if self._checkpoint_manager:
                        self._checkpoint_manager.record_error(error_msg)

                    example_results.append(
                        ExampleResult(
                            example_id=example.id,
                            strategy_name=self._strategy.name,
                            prediction="",
                            reference=example.reference,
                            correct=False,
                            score=0.0,
                            model_used="error",
                            error=str(e),
                            metadata=example.metadata,
                        )
                    )
                progress.increment()

        # Run all examples
        tasks = [process_one(ex) for ex in examples]
        await asyncio.gather(*tasks)

        # Log final rate limiter stats
        logger.info(f"Rate limiter stats: {self._rate_limiter.stats}")

        return example_results, errors

    async def _evaluate_single(self, example: BenchmarkExample) -> ExampleResult:
        """Evaluate a single example.

        Args:
            example: Benchmark example to evaluate.

        Returns:
            ExampleResult with prediction and scoring.
        """
        # Run inference
        strategy_result = await self._strategy.infer(
            example.query,
            system_prompt=self._system_prompt or example.system_prompt,
        )

        # Evaluate prediction
        eval_score = self._benchmark.evaluate(
            strategy_result.prediction,
            example,
        )

        # Build routing signals dict
        routing_signals = {}
        if strategy_result.routing:
            routing_signals = dict(strategy_result.routing.signals)

        return ExampleResult(
            example_id=example.id,
            strategy_name=self._strategy.name,
            prediction=strategy_result.prediction,
            reference=example.reference,
            correct=eval_score.correct,
            score=eval_score.score,
            model_used=strategy_result.model_used,
            tier=strategy_result.tier,
            latency_ms=strategy_result.latency_ms,
            cost_usd=strategy_result.cost_usd,
            tokens_in=strategy_result.tokens_in,
            tokens_out=strategy_result.tokens_out,
            metadata=example.metadata,
            routing_signals=routing_signals,
            error=strategy_result.error,
        )


async def run_evaluation(
    strategy: Strategy | str,
    benchmark: Benchmark | str,
    max_examples: int | None = None,
    system_prompt: str | None = None,
    subset: str | None = None,
    output_path: Path | str | None = None,
) -> EvaluationResult:
    """Run an evaluation and optionally save results.

    High-level convenience function for running evaluations.

    Args:
        strategy: Strategy instance or name (e.g., "tiered_router:eb1-gnn-preview-latest").
        benchmark: Benchmark instance or name (e.g., "gpqa_diamond").
        max_examples: Maximum examples to evaluate.
        system_prompt: Optional system prompt.
        subset: Benchmark subset.
        output_path: Optional path to save JSON results.

    Returns:
        EvaluationResult with all metrics.
    """
    # Resolve strategy
    if isinstance(strategy, str):
        strategy = _resolve_strategy(strategy)

    # Resolve benchmark
    if isinstance(benchmark, str):
        from embergraphrouter.evals.benchmarks import get_benchmark
        benchmark = get_benchmark(benchmark)

    # Create and run runner
    runner = SimpleRunner(
        strategy=strategy,
        benchmark=benchmark,
        max_examples=max_examples,
        system_prompt=system_prompt,
        subset=subset,
    )

    runner.setup()
    result = await runner.run()

    # Save results if output path provided
    if output_path:
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "w") as f:
            json.dump(result.to_dict(), f, indent=2)
        logger.info("Results saved to %s", output_path)

    return result


def _resolve_strategy(strategy_spec: str) -> Strategy:
    """Resolve strategy specification to Strategy instance.

    Formats:
    - "tiered_router" -> TieredRouterStrategy with defaults
    - "tiered_router:eb1-gnn-preview-12-27" -> TieredRouterStrategy with registry
    - "single_model:gpt-4" -> SingleModelStrategy with model
    - "random:model1,model2,model3" -> RandomModelStrategy with models

    Args:
        strategy_spec: Strategy specification string.

    Returns:
        Configured Strategy instance.
    """
    parts = strategy_spec.split(":", 1)
    strategy_type = parts[0]
    arg = parts[1] if len(parts) > 1 else None

    if strategy_type == "tiered_router":
        from embergraphrouter.evals.strategies import TieredRouterStrategy
        registry_name = arg or "eb1-gnn-preview-latest"
        return TieredRouterStrategy.from_registry(registry_name)

    elif strategy_type == "single_model":
        from embergraphrouter.evals.strategies import SingleModelStrategy
        if not arg:
            raise ValueError("single_model requires model name: single_model:gpt-4")
        return SingleModelStrategy(arg)

    elif strategy_type == "random":
        from embergraphrouter.evals.strategies import RandomModelStrategy
        if not arg:
            raise ValueError("random requires model list: random:model1,model2,model3")
        models = arg.split(",")
        return RandomModelStrategy(models)

    else:
        raise ValueError(f"Unknown strategy type: {strategy_type}")


__all__ = [
    "SimpleRunner",
    "run_evaluation",
]
