"""Campaign runner with smart checkpointing and adaptive parallelism.

This module provides infrastructure for running large-scale evaluation
campaigns across multiple benchmarks and strategies. Features:

- Persistent checkpointing for resume on interruption
- Adaptive parallelism based on rate limits and success rates
- Progress tracking and estimated time to completion
- Graceful degradation on errors
"""

import asyncio
import hashlib
import json
import logging
import pickle
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

from embergraphrouter.evals.benchmarks import Benchmark, get_benchmark
from embergraphrouter.evals.runners.base import (
    EvaluationResult,
    ExampleResult,
    Runner,
)
from embergraphrouter.evals.strategies.base import Strategy

logger = logging.getLogger(__name__)


@dataclass
class CheckpointState:
    """Persistent checkpoint state for campaign resumption.

    Attributes:
        campaign_id: Unique campaign identifier.
        started_at: Campaign start timestamp.
        updated_at: Last update timestamp.
        config_hash: Hash of campaign configuration.
        completed_tasks: Set of completed (strategy, benchmark, example_id).
        example_results: Accumulated results per (strategy, benchmark).
        errors: Accumulated errors.
        parallelism: Current parallelism level.
        rate_limit_hits: Count of rate limit errors.
    """

    campaign_id: str
    started_at: str
    updated_at: str
    config_hash: str
    completed_tasks: set[tuple[str, str, str]] = field(default_factory=set)
    example_results: dict[tuple[str, str], list[dict[str, Any]]] = field(
        default_factory=dict
    )
    errors: list[str] = field(default_factory=list)
    parallelism: int = 10
    rate_limit_hits: int = 0


@dataclass
class CampaignConfig:
    """Configuration for an evaluation campaign.

    Attributes:
        strategies: List of strategy names to evaluate.
        benchmarks: List of benchmark names to evaluate.
        max_examples_per_benchmark: Limit examples per benchmark (None = all).
        initial_parallelism: Starting concurrency level.
        min_parallelism: Minimum concurrency (for backoff).
        max_parallelism: Maximum concurrency.
        checkpoint_interval: Checkpoint every N examples.
        checkpoint_dir: Directory for checkpoint files.
        timeout_per_example: Timeout per example in seconds.
        adaptive_parallelism: Whether to adapt based on rate limits.
    """

    strategies: list[str]
    benchmarks: list[str]
    max_examples_per_benchmark: int | None = None
    initial_parallelism: int = 10
    min_parallelism: int = 1
    max_parallelism: int = 50
    checkpoint_interval: int = 100
    checkpoint_dir: str = "checkpoints"
    timeout_per_example: float = 120.0
    adaptive_parallelism: bool = True


@dataclass
class CampaignResult:
    """Result of a complete campaign.

    Attributes:
        campaign_id: Unique campaign identifier.
        started_at: Start timestamp.
        completed_at: Completion timestamp.
        config: Campaign configuration.
        results: Results per (strategy, benchmark).
        n_examples_total: Total examples processed.
        n_errors: Total errors encountered.
        duration_seconds: Total runtime.
    """

    campaign_id: str
    started_at: str
    completed_at: str
    config: CampaignConfig
    results: dict[str, dict[str, EvaluationResult]]  # strategy -> benchmark -> result
    n_examples_total: int
    n_errors: int
    duration_seconds: float

    def to_dict(self) -> dict[str, Any]:
        """Convert to serializable dictionary."""
        return {
            "campaign_id": self.campaign_id,
            "started_at": self.started_at,
            "completed_at": self.completed_at,
            "config": {
                "strategies": self.config.strategies,
                "benchmarks": self.config.benchmarks,
                "max_examples_per_benchmark": self.config.max_examples_per_benchmark,
                "initial_parallelism": self.config.initial_parallelism,
            },
            "n_examples_total": self.n_examples_total,
            "n_errors": self.n_errors,
            "duration_seconds": self.duration_seconds,
            "results": {
                strategy: {
                    benchmark: {
                        "accuracy": res.quality.get("accuracy", {}).get("mean", 0),
                        "n_examples": res.n_examples,
                    }
                    for benchmark, res in benchmarks.items()
                }
                for strategy, benchmarks in self.results.items()
            },
        }


class AdaptiveParallelism:
    """Manages adaptive parallelism based on success/failure rates.

    Reduces concurrency on rate limits, increases on success streaks.
    """

    def __init__(
        self,
        initial: int = 10,
        min_level: int = 1,
        max_level: int = 50,
        backoff_factor: float = 0.5,
        increase_factor: float = 1.2,
        success_threshold: int = 20,
    ):
        """Initialize adaptive parallelism controller.

        Args:
            initial: Starting concurrency level.
            min_level: Minimum concurrency.
            max_level: Maximum concurrency.
            backoff_factor: Multiply by this on rate limit.
            increase_factor: Multiply by this on success streak.
            success_threshold: Successes needed to increase.
        """
        self._current = initial
        self._min = min_level
        self._max = max_level
        self._backoff_factor = backoff_factor
        self._increase_factor = increase_factor
        self._success_threshold = success_threshold
        self._consecutive_successes = 0
        self._semaphore = asyncio.Semaphore(initial)

    @property
    def current(self) -> int:
        """Current parallelism level."""
        return self._current

    def record_success(self) -> None:
        """Record a successful request."""
        self._consecutive_successes += 1
        if self._consecutive_successes >= self._success_threshold:
            new_level = min(
                self._max,
                int(self._current * self._increase_factor),
            )
            if new_level > self._current:
                logger.info(
                    f"Adaptive parallelism: increasing from {self._current} to {new_level}"
                )
                self._current = new_level
                self._semaphore = asyncio.Semaphore(new_level)
            self._consecutive_successes = 0

    def record_rate_limit(self) -> None:
        """Record a rate limit error, reduce parallelism."""
        self._consecutive_successes = 0
        new_level = max(
            self._min,
            int(self._current * self._backoff_factor),
        )
        if new_level < self._current:
            logger.warning(
                f"Adaptive parallelism: reducing from {self._current} to {new_level}"
            )
            self._current = new_level
            self._semaphore = asyncio.Semaphore(new_level)

    def record_error(self) -> None:
        """Record a general error (doesn't affect parallelism much)."""
        self._consecutive_successes = 0

    async def acquire(self) -> None:
        """Acquire a slot for concurrent execution."""
        await self._semaphore.acquire()

    def release(self) -> None:
        """Release a slot."""
        self._semaphore.release()


class CampaignRunner(Runner):
    """Runner for large-scale evaluation campaigns.

    Features:
    - Persistent checkpointing for resume
    - Adaptive parallelism based on rate limits
    - Progress tracking
    - Graceful error handling

    Example:
        config = CampaignConfig(
            strategies=["tiered_router", "standalone:gpt-5-nano"],
            benchmarks=["gpqa_diamond", "simpleqa", "humaneval"],
            max_examples_per_benchmark=100,
        )
        runner = CampaignRunner(config)
        result = await runner.run()

    Resume interrupted campaign:
        runner = CampaignRunner.resume("checkpoints/campaign_abc123.pkl")
        result = await runner.run()
    """

    def __init__(
        self,
        config: CampaignConfig,
        strategies: dict[str, Strategy] | None = None,
        checkpoint_state: CheckpointState | None = None,
    ):
        """Initialize campaign runner.

        Args:
            config: Campaign configuration.
            strategies: Pre-constructed strategy instances.
            checkpoint_state: State to resume from.
        """
        self._config = config
        self._strategies = strategies or {}
        self._state = checkpoint_state
        self._parallelism: AdaptiveParallelism | None = None
        self._start_time: float = 0

    @classmethod
    def resume(cls, checkpoint_path: str | Path) -> "CampaignRunner":
        """Resume campaign from checkpoint file.

        Args:
            checkpoint_path: Path to checkpoint file.

        Returns:
            CampaignRunner configured to resume.

        Raises:
            FileNotFoundError: If checkpoint doesn't exist.
        """
        path = Path(checkpoint_path)
        if not path.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")

        with open(path, "rb") as f:
            state = pickle.load(f)

        # Reconstruct config from state
        # Note: In a full implementation, we'd store the full config
        config = CampaignConfig(
            strategies=[],  # Will be populated from state
            benchmarks=[],
            initial_parallelism=state.parallelism,
        )

        logger.info(
            f"Resuming campaign {state.campaign_id} "
            f"with {len(state.completed_tasks)} completed tasks"
        )

        return cls(config, checkpoint_state=state)

    def _generate_campaign_id(self) -> str:
        """Generate unique campaign ID."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        config_str = f"{self._config.strategies}{self._config.benchmarks}"
        hash_suffix = hashlib.sha256(config_str.encode()).hexdigest()[:8]
        return f"campaign_{timestamp}_{hash_suffix}"

    def _config_hash(self) -> str:
        """Generate hash of configuration for change detection."""
        config_str = json.dumps(
            {
                "strategies": sorted(self._config.strategies),
                "benchmarks": sorted(self._config.benchmarks),
                "max_examples": self._config.max_examples_per_benchmark,
            },
            sort_keys=True,
        )
        return hashlib.sha256(config_str.encode()).hexdigest()[:16]

    def _init_state(self) -> None:
        """Initialize or validate checkpoint state."""
        if self._state is None:
            self._state = CheckpointState(
                campaign_id=self._generate_campaign_id(),
                started_at=datetime.now().isoformat(),
                updated_at=datetime.now().isoformat(),
                config_hash=self._config_hash(),
                parallelism=self._config.initial_parallelism,
            )
        else:
            # Validate config hasn't changed
            current_hash = self._config_hash()
            if self._state.config_hash != current_hash:
                logger.warning(
                    f"Config changed since checkpoint. Old: {self._state.config_hash}, "
                    f"New: {current_hash}. Continuing with partial results."
                )

    def _checkpoint_path(self) -> Path:
        """Get checkpoint file path."""
        assert self._state is not None
        checkpoint_dir = Path(self._config.checkpoint_dir)
        checkpoint_dir.mkdir(parents=True, exist_ok=True)
        return checkpoint_dir / f"{self._state.campaign_id}.pkl"

    def _save_checkpoint(self) -> None:
        """Save current state to checkpoint file."""
        assert self._state is not None
        self._state.updated_at = datetime.now().isoformat()
        if self._parallelism:
            self._state.parallelism = self._parallelism.current

        path = self._checkpoint_path()
        with open(path, "wb") as f:
            pickle.dump(self._state, f)

        logger.debug(f"Checkpoint saved: {len(self._state.completed_tasks)} tasks")

    def _is_completed(
        self, strategy_name: str, benchmark_name: str, example_id: str
    ) -> bool:
        """Check if a task is already completed."""
        if self._state is None:
            return False
        return (strategy_name, benchmark_name, example_id) in self._state.completed_tasks

    def _mark_completed(
        self,
        strategy_name: str,
        benchmark_name: str,
        example_id: str,
        result: dict[str, Any],
    ) -> None:
        """Mark a task as completed and store result."""
        assert self._state is not None
        key = (strategy_name, benchmark_name, example_id)
        self._state.completed_tasks.add(key)

        results_key = (strategy_name, benchmark_name)
        if results_key not in self._state.example_results:
            self._state.example_results[results_key] = []
        self._state.example_results[results_key].append(result)

    async def _run_single_example(
        self,
        strategy: Strategy,
        strategy_name: str,
        benchmark: Benchmark,
        example: Any,  # BenchmarkExample
    ) -> ExampleResult | None:
        """Run a single example with error handling and checkpointing."""
        assert self._parallelism is not None
        assert self._state is not None

        example_id = example.id

        # Skip if already completed
        if self._is_completed(strategy_name, benchmark.name, example_id):
            return None

        await self._parallelism.acquire()
        try:
            # Run with timeout
            try:
                result = await asyncio.wait_for(
                    strategy.infer(
                        example.query,
                        system_prompt=example.system_prompt,
                    ),
                    timeout=self._config.timeout_per_example,
                )
            except asyncio.TimeoutError:
                self._parallelism.record_error()
                self._state.errors.append(f"Timeout: {example_id}")
                return ExampleResult(
                    example_id=example_id,
                    strategy_name=strategy_name,
                    prediction="",
                    reference=example.reference,
                    correct=False,
                    score=0.0,
                    model_used="timeout",
                    metadata={"error": "timeout"},
                    error="Timeout",
                )

            # Evaluate
            eval_score = benchmark.evaluate(result.prediction, example)

            # Check for rate limit in error
            if result.error and "rate" in result.error.lower():
                self._parallelism.record_rate_limit()
                self._state.rate_limit_hits += 1
            elif result.error:
                self._parallelism.record_error()
            else:
                self._parallelism.record_success()

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

            example_result = ExampleResult(
                example_id=example_id,
                strategy_name=strategy_name,
                prediction=result.prediction,
                reference=example.reference,
                correct=eval_score.correct,
                score=eval_score.score,
                model_used=result.model_used,
                tier=result.tier,
                latency_ms=result.latency_ms,
                cost_usd=result.cost_usd,
                tokens_in=result.tokens_in,
                tokens_out=result.tokens_out,
                metadata=example.metadata,
                routing_signals=routing_signals,
                error=result.error,
            )

            # Store in checkpoint
            self._mark_completed(
                strategy_name,
                benchmark.name,
                example_id,
                example_result.__dict__,
            )

            return example_result

        except Exception as e:
            logger.error(f"Error processing {example_id}: {e}", exc_info=True)
            self._parallelism.record_error()
            self._state.errors.append(f"Error {example_id}: {str(e)}")
            return ExampleResult(
                example_id=example_id,
                strategy_name=strategy_name,
                prediction="",
                reference=example.reference,
                correct=False,
                score=0.0,
                model_used="error",
                metadata={},
                error=str(e),
            )
        finally:
            self._parallelism.release()

    def _get_strategy(self, name: str) -> Strategy:
        """Get or create strategy by name."""
        if name in self._strategies:
            return self._strategies[name]

        # Parse strategy name
        if name.startswith("standalone:"):
            model = name.split(":", 1)[1]
            from embergraphrouter.evals.strategies import StandaloneStrategy

            strategy = StandaloneStrategy(model=model)
        elif name == "tiered_router":
            from embergraphrouter.evals.strategies import TieredRouterStrategy

            strategy = TieredRouterStrategy()
        else:
            raise ValueError(f"Unknown strategy: {name}")

        self._strategies[name] = strategy
        return strategy

    def _get_benchmark(self, name: str) -> Benchmark:
        """Get benchmark by name."""
        return get_benchmark(name)

    async def run(self) -> CampaignResult:
        """Run the full evaluation campaign.

        Returns:
            CampaignResult with all evaluation results.
        """
        self._init_state()
        assert self._state is not None

        self._start_time = time.perf_counter()

        # Initialize adaptive parallelism
        self._parallelism = AdaptiveParallelism(
            initial=self._state.parallelism,
            min_level=self._config.min_parallelism,
            max_level=self._config.max_parallelism,
        )

        logger.info(
            f"Starting campaign {self._state.campaign_id}: "
            f"{len(self._config.strategies)} strategies x "
            f"{len(self._config.benchmarks)} benchmarks"
        )

        all_results: dict[str, dict[str, EvaluationResult]] = {}
        n_examples_total = 0
        n_checkpoint = 0

        for strategy_name in self._config.strategies:
            strategy = self._get_strategy(strategy_name)
            all_results[strategy_name] = {}

            for benchmark_name in self._config.benchmarks:
                benchmark = self._get_benchmark(benchmark_name)

                logger.info(f"Running {strategy_name} on {benchmark_name}")

                # Collect examples
                examples = list(benchmark.load())
                if self._config.max_examples_per_benchmark:
                    examples = examples[: self._config.max_examples_per_benchmark]

                # Run examples
                tasks = [
                    self._run_single_example(strategy, strategy_name, benchmark, ex)
                    for ex in examples
                ]

                example_results: list[ExampleResult] = []
                for coro in asyncio.as_completed(tasks):
                    result = await coro
                    if result:
                        example_results.append(result)
                        n_examples_total += 1
                        n_checkpoint += 1

                        # Periodic checkpoint
                        if n_checkpoint >= self._config.checkpoint_interval:
                            self._save_checkpoint()
                            n_checkpoint = 0

                            # Log progress
                            elapsed = time.perf_counter() - self._start_time
                            rate = n_examples_total / elapsed if elapsed > 0 else 0
                            logger.info(
                                f"Progress: {n_examples_total} examples, "
                                f"{rate:.1f}/s, parallelism={self._parallelism.current}"
                            )

                # Aggregate results for this benchmark
                from embergraphrouter.evals.metrics import (
                    compute_efficiency_metrics,
                    compute_quality_metrics,
                    compute_routing_metrics,
                )

                quality = compute_quality_metrics(example_results)
                efficiency = compute_efficiency_metrics(example_results)
                routing = compute_routing_metrics(example_results)

                eval_result = EvaluationResult(
                    evaluation_id=self._state.campaign_id,
                    strategy_name=strategy_name,
                    benchmark_name=benchmark_name,
                    timestamp=datetime.now().isoformat(),
                    n_examples=len(example_results),
                    quality=quality,
                    efficiency=efficiency,
                    routing=routing,
                    by_difficulty={},
                    by_domain={},
                    by_tier={},
                    example_results=example_results,
                    errors=self._state.errors[-10:],  # Last 10 errors
                )

                all_results[strategy_name][benchmark_name] = eval_result

        # Final checkpoint
        self._save_checkpoint()

        duration = time.perf_counter() - self._start_time

        return CampaignResult(
            campaign_id=self._state.campaign_id,
            started_at=self._state.started_at,
            completed_at=datetime.now().isoformat(),
            config=self._config,
            results=all_results,
            n_examples_total=n_examples_total,
            n_errors=len(self._state.errors),
            duration_seconds=duration,
        )


async def run_campaign(
    strategies: list[str],
    benchmarks: list[str],
    max_examples: int | None = None,
    checkpoint_dir: str = "checkpoints",
    parallelism: int = 10,
) -> CampaignResult:
    """Convenience function to run a campaign.

    Args:
        strategies: Strategy names to evaluate.
        benchmarks: Benchmark names to evaluate.
        max_examples: Max examples per benchmark.
        checkpoint_dir: Directory for checkpoints.
        parallelism: Initial concurrency level.

    Returns:
        CampaignResult with all results.
    """
    config = CampaignConfig(
        strategies=strategies,
        benchmarks=benchmarks,
        max_examples_per_benchmark=max_examples,
        initial_parallelism=parallelism,
        checkpoint_dir=checkpoint_dir,
    )

    runner = CampaignRunner(config)
    return await runner.run()


__all__ = [
    "AdaptiveParallelism",
    "CampaignConfig",
    "CampaignResult",
    "CampaignRunner",
    "CheckpointState",
    "run_campaign",
]
