"""Shared checkpoint and resumption utilities for evaluation runners.

This module provides infrastructure for:
- Persistent checkpointing with atomic writes
- Graceful resumption after interruption
- Adaptive parallelism for rate limit handling
- Progress tracking with ETA estimation
"""

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

logger = logging.getLogger(__name__)

T = TypeVar("T")


@dataclass
class CheckpointData(Generic[T]):
    """Generic checkpoint data container.

    Attributes:
        run_id: Unique identifier for this run.
        started_at: ISO timestamp of start.
        updated_at: ISO timestamp of last update.
        config_hash: Hash of configuration for change detection.
        completed_ids: Set of completed task identifiers.
        results: Accumulated results (type depends on runner).
        errors: List of error messages.
        metadata: Additional run metadata.
    """

    run_id: str
    started_at: str
    updated_at: str
    config_hash: str
    completed_ids: set[str] = field(default_factory=set)
    results: dict[str, T] = field(default_factory=dict)
    errors: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)


class CheckpointManager:
    """Manages checkpoint persistence and resumption.

    Provides atomic checkpoint writes with backup, validation on resume,
    and progress tracking.

    Example:
        manager = CheckpointManager(
            checkpoint_dir="checkpoints",
            run_id="eval_20251228",
            config_hash="abc123",
        )

        # Check if task is done
        if not manager.is_completed("example_42"):
            result = await process(example)
            manager.mark_completed("example_42", result)

        # Periodic save
        if manager.should_save():
            manager.save()

        # Resume on restart
        manager = CheckpointManager.resume("checkpoints/eval_20251228.pkl")
    """

    def __init__(
        self,
        checkpoint_dir: str | Path,
        run_id: str | None = None,
        config_hash: str = "",
        save_interval: int = 100,
    ):
        """Initialize checkpoint manager.

        Args:
            checkpoint_dir: Directory for checkpoint files.
            run_id: Unique run identifier (generated if None).
            config_hash: Hash of configuration for change detection.
            save_interval: Save checkpoint every N completions.
        """
        self._checkpoint_dir = Path(checkpoint_dir)
        self._checkpoint_dir.mkdir(parents=True, exist_ok=True)

        self._run_id = run_id or self._generate_run_id()
        self._config_hash = config_hash
        self._save_interval = save_interval
        self._completions_since_save = 0

        self._data: CheckpointData[Any] = CheckpointData(
            run_id=self._run_id,
            started_at=datetime.now().isoformat(),
            updated_at=datetime.now().isoformat(),
            config_hash=config_hash,
        )

    @staticmethod
    def _generate_run_id() -> str:
        """Generate unique run ID."""
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        suffix = hashlib.sha256(str(time.time()).encode()).hexdigest()[:6]
        return f"run_{timestamp}_{suffix}"

    @property
    def run_id(self) -> str:
        """Get run identifier."""
        return self._run_id

    @property
    def checkpoint_path(self) -> Path:
        """Get checkpoint file path."""
        return self._checkpoint_dir / f"{self._run_id}.pkl"

    @property
    def n_completed(self) -> int:
        """Number of completed tasks."""
        return len(self._data.completed_ids)

    @property
    def errors(self) -> list[str]:
        """List of errors."""
        return self._data.errors

    def is_completed(self, task_id: str) -> bool:
        """Check if task is already completed."""
        return task_id in self._data.completed_ids

    def mark_completed(
        self,
        task_id: str,
        result: Any = None,
        auto_save: bool = True,
    ) -> None:
        """Mark a task as completed.

        Args:
            task_id: Unique task identifier.
            result: Optional result to store.
            auto_save: Whether to auto-save based on interval.
        """
        self._data.completed_ids.add(task_id)
        if result is not None:
            self._data.results[task_id] = result

        self._completions_since_save += 1

        if auto_save and self.should_save():
            self.save()

    def record_error(self, error: str) -> None:
        """Record an error message."""
        self._data.errors.append(error)

    def get_result(self, task_id: str) -> Any | None:
        """Get stored result for a task."""
        return self._data.results.get(task_id)

    def get_all_results(self) -> dict[str, Any]:
        """Get all stored results."""
        return dict(self._data.results)

    def should_save(self) -> bool:
        """Check if checkpoint should be saved based on interval."""
        return self._completions_since_save >= self._save_interval

    def save(self) -> None:
        """Save checkpoint atomically.

        Uses write-to-temp-then-rename for atomic updates.
        """
        self._data.updated_at = datetime.now().isoformat()

        # Write to temp file first
        temp_path = self.checkpoint_path.with_suffix(".tmp")
        with open(temp_path, "wb") as f:
            pickle.dump(self._data, f)

        # Atomic rename
        temp_path.rename(self.checkpoint_path)

        self._completions_since_save = 0
        logger.debug(
            f"Checkpoint saved: {self.n_completed} completed tasks"
        )

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

        Args:
            checkpoint_path: Path to checkpoint file.

        Returns:
            CheckpointManager with restored state.

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

        with open(path, "rb") as f:
            data: CheckpointData[Any] = pickle.load(f)

        manager = cls(
            checkpoint_dir=path.parent,
            run_id=data.run_id,
            config_hash=data.config_hash,
        )
        manager._data = data

        logger.info(
            f"Resumed checkpoint {data.run_id}: "
            f"{len(data.completed_ids)} completed tasks"
        )

        return manager

    def validate_config(self, expected_hash: str) -> bool:
        """Validate that config hasn't changed since checkpoint.

        Args:
            expected_hash: Expected config hash.

        Returns:
            True if configs match.
        """
        if self._data.config_hash != expected_hash:
            logger.warning(
                f"Config changed since checkpoint. "
                f"Old: {self._data.config_hash}, New: {expected_hash}"
            )
            return False
        return True

    def set_metadata(self, key: str, value: Any) -> None:
        """Set metadata value."""
        self._data.metadata[key] = value

    def get_metadata(self, key: str, default: Any = None) -> Any:
        """Get metadata value."""
        return self._data.metadata.get(key, default)


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

    Reduces concurrency on rate limits, increases on success streaks.
    Provides semaphore-based concurrency control.

    Example:
        limiter = AdaptiveRateLimiter(initial=10)

        async with limiter.acquire():
            try:
                result = await make_request()
                limiter.record_success()
            except RateLimitError:
                limiter.record_rate_limit()
    """

    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,
        cooldown_seconds: float = 1.0,
    ):
        """Initialize adaptive rate limiter.

        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.
            cooldown_seconds: Wait time after rate limit before retrying.
        """
        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._cooldown = cooldown_seconds

        self._consecutive_successes = 0
        self._total_successes = 0
        self._total_rate_limits = 0
        self._total_errors = 0

        self._semaphore = asyncio.Semaphore(initial)
        self._lock = asyncio.Lock()

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

    @property
    def stats(self) -> dict[str, int]:
        """Get rate limiter statistics."""
        return {
            "current_level": self._current,
            "total_successes": self._total_successes,
            "total_rate_limits": self._total_rate_limits,
            "total_errors": self._total_errors,
        }

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

        Returns:
            Context manager that releases on exit.
        """
        await self._semaphore.acquire()
        return RateLimiterContext(self)

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

    async def record_success(self) -> None:
        """Record a successful request."""
        async with self._lock:
            self._consecutive_successes += 1
            self._total_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"Rate limiter: increasing parallelism "
                        f"from {self._current} to {new_level}"
                    )
                    self._current = new_level
                    # Recreate semaphore with new limit
                    self._semaphore = asyncio.Semaphore(new_level)
                self._consecutive_successes = 0

    async def record_rate_limit(self) -> None:
        """Record a rate limit error, reduce parallelism."""
        async with self._lock:
            self._consecutive_successes = 0
            self._total_rate_limits += 1

            new_level = max(
                self._min,
                int(self._current * self._backoff_factor),
            )
            if new_level < self._current:
                logger.warning(
                    f"Rate limiter: reducing parallelism "
                    f"from {self._current} to {new_level}"
                )
                self._current = new_level
                self._semaphore = asyncio.Semaphore(new_level)

        # Cooldown before allowing more requests
        await asyncio.sleep(self._cooldown)

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


class RateLimiterContext:
    """Context manager for rate limiter slot."""

    def __init__(self, limiter: AdaptiveRateLimiter):
        self._limiter = limiter

    async def __aenter__(self) -> "RateLimiterContext":
        return self

    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        self._limiter.release()


@dataclass
class ProgressTracker:
    """Tracks evaluation progress with ETA estimation.

    Attributes:
        total: Total number of tasks.
        completed: Number completed.
        start_time: Start timestamp.
    """

    total: int
    completed: int = 0
    start_time: float = field(default_factory=time.time)

    @property
    def elapsed_seconds(self) -> float:
        """Elapsed time in seconds."""
        return time.time() - self.start_time

    @property
    def rate(self) -> float:
        """Tasks per second."""
        elapsed = self.elapsed_seconds
        return self.completed / elapsed if elapsed > 0 else 0.0

    @property
    def eta_seconds(self) -> float:
        """Estimated time remaining in seconds."""
        remaining = self.total - self.completed
        return remaining / self.rate if self.rate > 0 else float("inf")

    @property
    def progress_pct(self) -> float:
        """Progress as percentage."""
        return 100 * self.completed / self.total if self.total > 0 else 0.0

    def increment(self, n: int = 1) -> None:
        """Increment completed count."""
        self.completed += n

    def format_status(self) -> str:
        """Format progress status string."""
        eta = self.eta_seconds
        if eta == float("inf"):
            eta_str = "calculating..."
        elif eta > 3600:
            eta_str = f"{eta / 3600:.1f}h"
        elif eta > 60:
            eta_str = f"{eta / 60:.1f}m"
        else:
            eta_str = f"{eta:.0f}s"

        return (
            f"Progress: {self.completed}/{self.total} ({self.progress_pct:.1f}%) "
            f"| Rate: {self.rate:.1f}/s | ETA: {eta_str}"
        )


def compute_config_hash(config: dict[str, Any]) -> str:
    """Compute deterministic hash of configuration.

    Args:
        config: Configuration dictionary.

    Returns:
        16-character hash string.
    """
    import json

    # Sort keys for determinism
    config_str = json.dumps(config, sort_keys=True, default=str)
    return hashlib.sha256(config_str.encode()).hexdigest()[:16]


__all__ = [
    "AdaptiveRateLimiter",
    "CheckpointData",
    "CheckpointManager",
    "ProgressTracker",
    "RateLimiterContext",
    "compute_config_hash",
]
