"""MIPROv2 optimization orchestration."""

import logging
from typing import List, Dict, Any, Tuple

import dspy
from dspy.evaluate import Evaluate
from dspy.teleprompt import MIPROv2

from .training_data import TrainingExample
from .pipeline import VideoAnalysisPipeline
from .metrics import VideoAnalysisMetric

logger = logging.getLogger(__name__)


class PromptTuner:
    """Orchestrates MIPROv2 optimization of the two video-analyzer prompts.

    Args:
        lm_config: LLM configuration dict with keys: type, model, api_base, api_key
        num_candidates: Number of prompt variations MIPROv2 generates per module
        num_trials: Number of Bayesian optimization trials MIPROv2 runs
        max_bootstrapped_demos: Max few-shot demos generated by bootstrapping
        max_labeled_demos: Max few-shot demos taken directly from training data
        description_weight: Weight for final description score in metric (0.0-1.0)
    """

    def __init__(
        self,
        lm_config: Dict[str, Any],
        num_candidates: int = 10,
        num_trials: int = 20,
        max_bootstrapped_demos: int = 3,
        max_labeled_demos: int = 4,
        description_weight: float = 0.7,
    ):
        self.lm_config = lm_config
        self.num_candidates = num_candidates
        self.num_trials = num_trials
        self.max_bootstrapped_demos = max_bootstrapped_demos
        self.max_labeled_demos = max_labeled_demos
        self.description_weight = description_weight

    def _configure_lm(self) -> None:
        """Configure the DSPy global LM from lm_config."""
        client_type = self.lm_config.get("type", "ollama")
        model = self.lm_config.get("model", "llama3.2-vision")

        if client_type == "ollama":
            lm = dspy.LM(
                model=f"ollama/{model}",
                api_base=self.lm_config.get("api_base", "http://localhost:11434"),
                api_key="ollama",
            )
        else:  # openai_api — prefix with openai/ so LiteLLM routes correctly
            lm = dspy.LM(
                model=f"openai/{model}",
                api_base=self.lm_config.get("api_base"),
                api_key=self.lm_config.get("api_key"),
            )

        dspy.configure(lm=lm)
        logger.info(f"Configured DSPy LM: {model} ({client_type})")

    def _build_dspy_examples(self, training_examples: List[TrainingExample]) -> List[dspy.Example]:
        """Convert TrainingExample objects to dspy.Example objects."""
        dspy_examples = []

        for te in training_examples:
            frames = []
            for tf in te.frames:
                frame_dict: Dict[str, Any] = {
                    "index": tf.index,
                    "timestamp": tf.timestamp,
                    "image": None,
                }
                if tf.image_path and tf.image_path.exists():
                    try:
                        frame_dict["image"] = dspy.Image.from_file(str(tf.image_path))
                    except Exception as e:
                        logger.warning(f"Could not load image {tf.image_path}: {e}")
                frames.append(frame_dict)

            # Match VideoAnalyzer._format_user_prompt() exactly
            user_question = f"I want to know {te.user_question}" if te.user_question else ""

            ideal_frame_notes = [tf.ideal_note or "" for tf in te.frames]

            ex = dspy.Example(
                frames=frames,
                user_question=user_question,
                transcript=te.transcript,
                ideal_description=te.ideal_description,
                frame_notes_list=ideal_frame_notes,
                has_ideal_frame_notes=te.has_ideal_frame_notes,
            ).with_inputs("frames", "user_question", "transcript")

            dspy_examples.append(ex)

        return dspy_examples

    def _split_examples(
        self, examples: List[dspy.Example]
    ) -> Tuple[List[dspy.Example], List[dspy.Example]]:
        """Split examples into train and validation sets.

        Uses 80/20 split when >= 3 examples, otherwise reuses training set as
        validation (common when users have only 1-2 videos to tune with).
        """
        if len(examples) >= 3:
            split_idx = max(1, int(len(examples) * 0.8))
            return examples[:split_idx], examples[split_idx:]
        return examples, examples

    def optimize(self, training_examples: List[TrainingExample]) -> dspy.Module:
        """Run MIPROv2 optimization and return the optimized pipeline.

        Args:
            training_examples: Training examples loaded from edited analysis.json files

        Returns:
            Optimized VideoAnalysisPipeline with improved prompt instructions
        """
        self._configure_lm()

        pipeline = VideoAnalysisPipeline()
        metric = VideoAnalysisMetric(description_weight=self.description_weight)

        dspy_examples = self._build_dspy_examples(training_examples)
        trainset, valset = self._split_examples(dspy_examples)

        logger.info(
            f"Starting optimization: {len(trainset)} train / {len(valset)} val examples, "
            f"{self.num_candidates} candidates, {self.num_trials} trials"
        )

        evaluator = Evaluate(
            devset=valset,
            metric=metric,
            num_threads=1,
            display_progress=False,
            display_table=False,
        )

        logger.info("Evaluating baseline (unoptimized) prompts...")
        baseline_score = float(evaluator(pipeline))
        logger.info(f"Baseline score: {baseline_score:.1f}%")

        optimizer = MIPROv2(
            metric=metric,
            auto=None,
            num_candidates=self.num_candidates,
            max_bootstrapped_demos=self.max_bootstrapped_demos,
            max_labeled_demos=self.max_labeled_demos,
            verbose=True,
        )

        optimized_pipeline = optimizer.compile(
            pipeline,
            trainset=trainset,
            valset=valset,
            num_trials=self.num_trials,
            minibatch=False,
        )

        logger.info("Evaluating optimized prompts...")
        optimized_score = float(evaluator(optimized_pipeline))

        improvement = optimized_score - baseline_score
        logger.info("")
        logger.info("=" * 40)
        logger.info(f"  Baseline score:  {baseline_score:.1f}%")
        logger.info(f"  Optimized score: {optimized_score:.1f}%")
        logger.info(f"  Improvement:     {improvement:+.1f}%")
        logger.info("=" * 40)
        logger.info("")

        return optimized_pipeline
