"""
Optimization scheduling and metrics tracking.

Provides automatic background optimization that triggers based on
trace thresholds or time intervals, following the ConsolidationJob pattern.

Requirements: SEC-08
"""

import asyncio
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional, List, Dict, Any
from pathlib import Path
import json
import time
import uuid

# Optional DSPy imports
try:
    import dspy

    HAS_DSPY = True
except ImportError:
    HAS_DSPY = False
    dspy = None

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from gsd_rlm.optimization.traces import TraceCollector
    from gsd_rlm.optimization.optimizer import AgentOptimizer


def _utc_now_iso() -> str:
    """Get current UTC time as ISO format string."""
    return datetime.now(timezone.utc).isoformat()


@dataclass
class OptimizationRun:
    """
    Record of a single optimization run.

    Captures the before/after scores, improvement metrics, and
    execution details for tracking optimization history.
    """

    run_id: str
    timestamp: str
    traces_used: int
    score_before: float
    score_after: float
    improvement: float
    success: bool
    duration_ms: int = 0
    error: Optional[str] = None

    def to_dict(self) -> dict:
        """
        Convert run to dictionary for serialization.

        Returns:
            Dictionary representation suitable for JSON storage.
        """
        return {
            "run_id": self.run_id,
            "timestamp": self.timestamp,
            "traces_used": self.traces_used,
            "score_before": self.score_before,
            "score_after": self.score_after,
            "improvement": self.improvement,
            "success": self.success,
            "duration_ms": self.duration_ms,
            "error": self.error,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "OptimizationRun":
        """
        Create OptimizationRun from dictionary.

        Args:
            data: Dictionary with run fields.

        Returns:
            OptimizationRun instance.
        """
        return cls(
            run_id=data.get("run_id", ""),
            timestamp=data.get("timestamp", _utc_now_iso()),
            traces_used=data.get("traces_used", 0),
            score_before=data.get("score_before", 0.0),
            score_after=data.get("score_after", 0.0),
            improvement=data.get("improvement", 0.0),
            success=data.get("success", False),
            duration_ms=data.get("duration_ms", 0),
            error=data.get("error"),
        )


class OptimizationMetricsTracker:
    """
    Track optimization run history and aggregate metrics.

    Provides:
    - History of all optimization runs
    - Aggregate statistics (avg improvement, success rate)
    - Persistence to JSON file

    Usage:
        tracker = OptimizationMetricsTracker("metrics/optimization.json")
        tracker.record(run)
        avg_improvement = tracker.get_average_improvement()
    """

    def __init__(self, metrics_path: Optional[str] = None):
        """
        Initialize the metrics tracker.

        Args:
            metrics_path: Optional path to persist metrics to JSON.
        """
        self.metrics_path = metrics_path
        self.runs: List[OptimizationRun] = []
        if metrics_path:
            self._load()

    def record(self, run: OptimizationRun) -> None:
        """
        Record an optimization run.

        Args:
            run: OptimizationRun to record.
        """
        self.runs.append(run)
        if self.metrics_path:
            self._save()

    def get_latest(self) -> Optional[OptimizationRun]:
        """
        Get the most recent optimization run.

        Returns:
            Latest OptimizationRun or None if no runs recorded.
        """
        return self.runs[-1] if self.runs else None

    def get_average_improvement(self, last_n: int = 10) -> float:
        """
        Get average improvement from last N runs.

        Only considers successful runs for improvement calculation.

        Args:
            last_n: Number of recent runs to consider.

        Returns:
            Average improvement score, or 0.0 if no successful runs.
        """
        if not self.runs:
            return 0.0

        recent = self.runs[-last_n:]
        improvements = [r.improvement for r in recent if r.success]
        return sum(improvements) / len(improvements) if improvements else 0.0

    def get_success_rate(self, last_n: int = 10) -> float:
        """
        Get success rate from last N runs.

        Args:
            last_n: Number of recent runs to consider.

        Returns:
            Success rate as a float between 0.0 and 1.0.
        """
        if not self.runs:
            return 0.0

        recent = self.runs[-last_n:]
        successes = sum(1 for r in recent if r.success)
        return successes / len(recent)

    def get_total_runs(self) -> int:
        """
        Get total number of optimization runs.

        Returns:
            Total count of recorded runs.
        """
        return len(self.runs)

    def _save(self) -> None:
        """Save metrics to JSON file."""
        if not self.metrics_path:
            return

        Path(self.metrics_path).parent.mkdir(parents=True, exist_ok=True)
        data = {"runs": [r.to_dict() for r in self.runs]}
        with open(self.metrics_path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)

    def _load(self) -> None:
        """Load metrics from JSON file."""
        if not self.metrics_path or not Path(self.metrics_path).exists():
            return

        try:
            with open(self.metrics_path, "r", encoding="utf-8") as f:
                data = json.load(f)
            self.runs = [
                OptimizationRun.from_dict(run_data) for run_data in data.get("runs", [])
            ]
        except (json.JSONDecodeError, KeyError, TypeError):
            self.runs = []


class OptimizationScheduler:
    """
    Schedule automatic optimization runs.

    Triggers optimization based on:
    - Trace count threshold (unoptimized traces)
    - Time interval (periodic optimization)
    - Manual trigger (run_once)

    Follows ConsolidationJob pattern for consistency.

    Usage:
        scheduler = OptimizationScheduler(
            trace_collector=collector,
            optimizer=optimizer,
            interval_seconds=3600,  # 1 hour
            trace_threshold=50,
        )

        await scheduler.start()  # Start background loop

        # Or run manually
        result = await scheduler.run_once()
    """

    DEFAULT_TRACE_THRESHOLD = 50
    DEFAULT_INTERVAL_SECONDS = 3600  # 1 hour
    MIN_TRAINSET_SIZE = 10

    def __init__(
        self,
        trace_collector: "TraceCollector",
        optimizer: "AgentOptimizer",
        metrics_tracker: Optional[OptimizationMetricsTracker] = None,
        interval_seconds: int = DEFAULT_INTERVAL_SECONDS,
        trace_threshold: int = DEFAULT_TRACE_THRESHOLD,
    ):
        """
        Initialize the optimization scheduler.

        Args:
            trace_collector: TraceCollector to get training data from.
            optimizer: AgentOptimizer to run optimizations.
            metrics_tracker: Optional tracker for optimization history.
            interval_seconds: Interval between optimization checks.
            trace_threshold: Minimum traces to trigger optimization.
        """
        self.trace_collector = trace_collector
        self.optimizer = optimizer
        self.metrics_tracker = metrics_tracker or OptimizationMetricsTracker()
        self.interval_seconds = interval_seconds
        self.trace_threshold = trace_threshold

        # Scheduler state
        self._running = False
        self._task: Optional[asyncio.Task] = None
        self._last_run: Optional[str] = None

    async def start(self) -> None:
        """Start the background optimization loop."""
        if self._running:
            return

        self._running = True
        self._task = asyncio.create_task(self._optimization_loop())

    async def stop(self) -> None:
        """Stop the background optimization loop."""
        self._running = False

        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None

    async def run_once(self) -> Optional[OptimizationRun]:
        """
        Run optimization once if threshold met.

        Returns:
            OptimizationRun if optimization was performed, None otherwise
        """
        if not self._should_optimize():
            return None
        return await self._do_optimization()

    async def force_run(self) -> Optional[OptimizationRun]:
        """
        Force optimization run regardless of threshold.

        Returns:
            OptimizationRun with optimization results
        """
        return await self._do_optimization()

    def _should_optimize(self) -> bool:
        """
        Check if optimization should be triggered.

        Returns:
            True if DSPy is available and trace threshold is met.
        """
        if not HAS_DSPY:
            return False

        # Check trace count
        successful_traces = self.trace_collector.get_successful_traces(
            limit=self.trace_threshold + 1
        )
        return len(successful_traces) >= self.trace_threshold

    async def _optimization_loop(self) -> None:
        """Main background optimization loop."""
        while self._running:
            try:
                if self._should_optimize():
                    await self._do_optimization()

                await asyncio.sleep(self.interval_seconds)

            except asyncio.CancelledError:
                break
            except Exception:
                # Log error but continue running
                await asyncio.sleep(self.interval_seconds)

    async def _do_optimization(self) -> OptimizationRun:
        """
        Perform optimization run.

        Returns:
            OptimizationRun with results or error information.
        """
        start_time = time.time()
        run_id = f"opt_{uuid.uuid4().hex[:12]}"

        try:
            # Get training data
            trainset = self.trace_collector.get_training_dataset()

            if len(trainset) < self.MIN_TRAINSET_SIZE:
                return OptimizationRun(
                    run_id=run_id,
                    timestamp=_utc_now_iso(),
                    traces_used=len(trainset),
                    score_before=0.0,
                    score_after=0.0,
                    improvement=0.0,
                    success=False,
                    error=f"Insufficient training data (need at least {self.MIN_TRAINSET_SIZE} traces, got {len(trainset)})",
                )

            # Run optimization (synchronous call to optimizer)
            result = self.optimizer.optimize(trainset)

            # Calculate duration
            duration_ms = int((time.time() - start_time) * 1000)

            # Create run record
            run = OptimizationRun(
                run_id=run_id,
                timestamp=_utc_now_iso(),
                traces_used=len(trainset),
                score_before=result.score_before,
                score_after=result.score_after,
                improvement=result.improvement,
                success=result.success,
                duration_ms=duration_ms,
                error=result.error,
            )

            # Record metrics
            self.metrics_tracker.record(run)
            self._last_run = run.timestamp

            return run

        except Exception as e:
            duration_ms = int((time.time() - start_time) * 1000)
            run = OptimizationRun(
                run_id=run_id,
                timestamp=_utc_now_iso(),
                traces_used=0,
                score_before=0.0,
                score_after=0.0,
                improvement=0.0,
                success=False,
                duration_ms=duration_ms,
                error=str(e),
            )
            self.metrics_tracker.record(run)
            return run

    @property
    def is_running(self) -> bool:
        """Check if the scheduler is currently running."""
        return self._running

    @property
    def last_run(self) -> Optional[str]:
        """Get ISO timestamp of last optimization."""
        return self._last_run

    @property
    def metrics(self) -> OptimizationMetricsTracker:
        """Get the metrics tracker."""
        return self.metrics_tracker
