"""
DSPy MIPROv2 optimization wrapper for agent prompt improvement.

Provides AgentOptimizer class that wraps DSPy's MIPROv2 optimizer
to take execution traces as training data and produce optimized prompts.

Requirements: SEC-06
"""

import copy
import json
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, List, Callable, Any, Dict

try:
    import dspy
    from dspy import MIPROv2

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

from gsd_rlm.optimization.metrics import task_success_metric


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


@dataclass
class OptimizationResult:
    """
    Result of an optimization run.

    Captures the before/after scores, improvement metrics, and
    configuration used for the optimization.
    """

    success: bool
    score_before: float
    score_after: float
    improvement: float
    traces_used: int
    timestamp: str = field(default_factory=_utc_now_iso)
    optimizer_config: Dict[str, Any] = field(default_factory=dict)
    optimized_program: Optional[Any] = None
    error: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        """
        Convert result to dictionary for serialization.

        Returns:
            Dictionary representation (excludes optimized_program).
        """
        return {
            "success": self.success,
            "score_before": self.score_before,
            "score_after": self.score_after,
            "improvement": self.improvement,
            "traces_used": self.traces_used,
            "timestamp": self.timestamp,
            "optimizer_config": self.optimizer_config,
            "error": self.error,
        }

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

        Args:
            data: Dictionary with result fields.

        Returns:
            OptimizationResult instance.
        """
        return cls(
            success=data.get("success", False),
            score_before=data.get("score_before", 0.0),
            score_after=data.get("score_after", 0.0),
            improvement=data.get("improvement", 0.0),
            traces_used=data.get("traces_used", 0),
            timestamp=data.get("timestamp", _utc_now_iso()),
            optimizer_config=data.get("optimizer_config", {}),
            error=data.get("error"),
        )


class AgentOptimizer:
    """
    Wrapper for DSPy MIPROv2 prompt optimization.

    Provides a simplified interface for optimizing agent prompts using
    DSPy's MIPROv2 algorithm. Takes execution traces as training data
    and produces optimized prompt configurations.

    Usage:
        optimizer = AgentOptimizer(agent_module=my_agent)
        result = optimizer.optimize(trainset=training_examples)
        if result.success:
            optimizer.save(result.optimized_program, Path("optimized.json"))
    """

    DEFAULT_MAX_ROUNDS = 3
    DEFAULT_THRESHOLD = 0.9
    MIN_TRAINSET_SIZE = 10

    def __init__(
        self,
        agent_module: Optional[Any] = None,
        metric: Callable = task_success_metric,
        auto: str = "light",
        prompt_model: Optional[Any] = None,
        max_bootstrapped_demos: int = 4,
        max_labeled_demos: int = 16,
        max_rounds: int = DEFAULT_MAX_ROUNDS,
        num_threads: int = 1,
    ):
        """
        Initialize the AgentOptimizer.

        Args:
            agent_module: DSPy module to optimize (e.g., ReAct, ChainOfThought).
            metric: Evaluation metric function. Defaults to task_success_metric.
            auto: Optimization effort level: "light", "medium", or "heavy".
            prompt_model: Model for prompt generation (defaults to configured LM).
            max_bootstrapped_demos: Max demos to bootstrap during optimization.
            max_labeled_demos: Max labeled demos to include.
            max_rounds: Max optimization rounds.
            num_threads: Number of parallel threads for evaluation.

        Raises:
            ImportError: If DSPy is not installed.
        """
        if not HAS_DSPY:
            raise ImportError(
                "DSPy is required for AgentOptimizer. "
                "Install it with: pip install dspy-ai"
            )

        self.agent_module = agent_module
        self.metric = metric
        self.auto = auto
        self.prompt_model = prompt_model
        self.max_bootstrapped_demos = max_bootstrapped_demos
        self.max_labeled_demos = max_labeled_demos
        self.max_rounds = max_rounds
        self.num_threads = num_threads

        # Store optimization history
        self._history: List[OptimizationResult] = []

    def optimize(
        self,
        trainset: List[Any],
        valset: Optional[List[Any]] = None,
        num_trials: Optional[int] = None,
    ) -> OptimizationResult:
        """
        Run MIPROv2 optimization on agent module.

        Args:
            trainset: List of DSPy Examples for training.
            valset: Optional validation set. If not provided, uses trainset.
            num_trials: Number of optimization trials. Auto-set based on `auto` mode.

        Returns:
            OptimizationResult with scores and optimized program.

        Raises:
            ValueError: If trainset has fewer than MIN_TRAINSET_SIZE examples.
            RuntimeError: If optimization fails.
        """
        # Validate trainset size
        if len(trainset) < self.MIN_TRAINSET_SIZE:
            raise ValueError(
                f"Training set must have at least {self.MIN_TRAINSET_SIZE} examples, "
                f"got {len(trainset)}. Collect more execution traces."
            )

        if self.agent_module is None:
            raise ValueError(
                "agent_module is required for optimization. "
                "Provide it in __init__ or set self.agent_module."
            )

        # Set num_trials based on auto mode if not specified
        if num_trials is None:
            auto_trials = {"light": 10, "medium": 25, "heavy": 50}
            num_trials = auto_trials.get(self.auto, 10)

        # Use trainset as valset if not provided
        evalset = valset if valset is not None else trainset

        try:
            # Calculate score before optimization
            score_before = self._evaluate(self.agent_module, evalset)

            # Create MIPROv2 optimizer with configured parameters
            mipro = MIPROv2(
                metric=self.metric,
                prompt_model=self.prompt_model,
                auto=self.auto,
                num_threads=self.num_threads,
                max_bootstrapped_demos=self.max_bootstrapped_demos,
                max_labeled_demos=self.max_labeled_demos,
            )

            # Deep copy the module to avoid modifying the original
            module_copy = copy.deepcopy(self.agent_module)

            # Run optimization
            optimized_program = mipro.compile(
                module_copy,
                trainset=trainset,
                valset=evalset,
                num_trials=num_trials,
                max_rounds=self.max_rounds,
            )

            # Calculate score after optimization
            score_after = self._evaluate(optimized_program, evalset)

            # Calculate improvement
            improvement = score_after - score_before

            # Build result
            result = OptimizationResult(
                success=True,
                score_before=score_before,
                score_after=score_after,
                improvement=improvement,
                traces_used=len(trainset),
                optimizer_config={
                    "auto": self.auto,
                    "num_trials": num_trials,
                    "max_rounds": self.max_rounds,
                    "max_bootstrapped_demos": self.max_bootstrapped_demos,
                    "max_labeled_demos": self.max_labeled_demos,
                },
                optimized_program=optimized_program,
            )

            # Store in history
            self._history.append(result)

            return result

        except Exception as e:
            # Return error result
            result = OptimizationResult(
                success=False,
                score_before=0.0,
                score_after=0.0,
                improvement=0.0,
                traces_used=len(trainset),
                error=str(e),
            )
            self._history.append(result)
            return result

    def _evaluate(self, program: Any, dataset: List[Any]) -> float:
        """
        Evaluate a program on a dataset using the configured metric.

        Args:
            program: DSPy program to evaluate.
            dataset: List of examples to evaluate on.

        Returns:
            Average score across the dataset.
        """
        if not dataset:
            return 0.0

        total_score = 0.0
        for example in dataset:
            try:
                # Get prediction from program
                # Handle different input formats
                if hasattr(example, "task_input"):
                    pred = program(task_input=example.task_input)
                elif hasattr(example, "question"):
                    pred = program(question=example.question)
                else:
                    # Try common input field names
                    pred = program(**example.inputs())

                # Evaluate with metric
                score = self.metric(example, pred)
                total_score += score
            except Exception:
                # Failed prediction gets 0 score
                total_score += 0.0

        return total_score / len(dataset)

    def save(self, program: Any, path: Path) -> None:
        """
        Save optimized program to file.

        Args:
            program: Optimized DSPy program to save.
            path: Path to save the program (JSON format).
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        if hasattr(program, "save"):
            # Use DSPy's built-in save if available
            program.save(str(path))
        else:
            # Fallback: save as JSON with state info
            state = {
                "type": type(program).__name__,
                "timestamp": _utc_now_iso(),
            }

            # Try to extract signature information
            if hasattr(program, "signature"):
                state["signature"] = str(program.signature)

            # Save demonstrations if available
            if hasattr(program, "demos"):
                state["demos"] = [str(demo) for demo in program.demos]

            with open(path, "w", encoding="utf-8") as f:
                json.dump(state, f, indent=2, default=str)

    def load(self, path: Path) -> Optional[Any]:
        """
        Load optimized program from file.

        Args:
            path: Path to the saved program.

        Returns:
            Loaded DSPy program, or None if loading fails.
        """
        if not HAS_DSPY:
            return None

        path = Path(path)
        if not path.exists():
            return None

        try:
            # Check if it's a DSPy save format
            with open(path, "r", encoding="utf-8") as f:
                content = f.read()

            # Try to parse as JSON first
            try:
                data = json.loads(content)
                # This is our JSON format, not DSPy's native format
                # Return the data dict for manual reconstruction
                return data
            except json.JSONDecodeError:
                # Not JSON, might be DSPy's native format
                pass

            # Try DSPy's load if available on the module
            if self.agent_module and hasattr(self.agent_module, "load"):
                return self.agent_module.load(str(path))

            return None

        except Exception:
            return None

    def get_history(self) -> List[OptimizationResult]:
        """
        Get optimization history.

        Returns:
            List of all OptimizationResult from this optimizer.
        """
        return self._history.copy()

    def clear_history(self) -> None:
        """Clear optimization history."""
        self._history.clear()

    @property
    def last_result(self) -> Optional[OptimizationResult]:
        """Get the most recent optimization result."""
        return self._history[-1] if self._history else None
