"""Benchmark engine orchestrating intent elaboration, parallel runs, judging, and ranking."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

import litellm

from sage_evaluator.benchmark.collector import MetricsCollector
from sage_evaluator.benchmark.runner import BenchmarkRunner
from sage_evaluator.discovery.pricing import PricingLookup
from sage_evaluator.evaluation.judge import LLMJudge
from sage_evaluator.evaluation.rubrics import load_rubric
from sage_evaluator.exceptions import BenchmarkError
from sage_evaluator.models import (
    BenchmarkIntent,
    BenchmarkReport,
    EvaluationRubric,
    ModelBenchmarkResult,
    ModelRunResult,
    QualityScore,
)

if TYPE_CHECKING:
    from sage.main_config import MainConfig

logger = logging.getLogger(__name__)

INTENT_ELABORATION_PROMPT = """\
You are helping prepare a benchmark evaluation for an AI agent.

Given the user's intent below, produce a structured elaboration with:
1. A clearer, more detailed version of the intent
2. What a good output would look like (expected outcome)
3. Specific evaluation criteria (3-5 bullet points)

User's intent: {intent}

Respond with JSON:
{{
  "elaborated_intent": "...",
  "expected_outcome": "...",
  "evaluation_criteria": ["...", "..."]
}}
"""


class BenchmarkEngine:
    """Orchestrates the full benchmark pipeline:
    1. Elaborate user intent via evaluator model
    2. Run agent against each model with InstrumentedProvider
    3. Score each output via LLM-as-judge (optional)
    4. Compute costs from pricing data
    5. Rank models by composite score
    """

    def __init__(
        self,
        evaluator_model: str,
        rubric: str | EvaluationRubric = "default",
        runs: int = 1,
        use_judge: bool = True,
        quality_weight: float = 0.5,
        cost_weight: float = 0.25,
        speed_weight: float = 0.25,
        main_config: MainConfig | None = None,
    ):
        self.evaluator_model = evaluator_model
        self.rubric = rubric if isinstance(rubric, EvaluationRubric) else load_rubric(rubric)
        self.runs = max(1, runs)
        self.use_judge = use_judge
        self.quality_weight = quality_weight
        self.cost_weight = cost_weight
        self.speed_weight = speed_weight
        self.main_config = main_config
        self.runner = BenchmarkRunner()
        self.judge = LLMJudge(model=evaluator_model) if use_judge else None
        self.pricing = PricingLookup()

    async def run(
        self,
        config_path: str,
        models: list[str],
        intent: str,
        model_params: dict[str, Any] | None = None,
    ) -> BenchmarkReport:
        """Run the full benchmark pipeline.

        Args:
            config_path: Path to agent config.
            models: List of model identifiers to benchmark.
            intent: User intent to benchmark against.
            model_params: Optional model parameters.

        Returns:
            BenchmarkReport with ranked results.
        """
        if not models:
            raise BenchmarkError("No models specified for benchmarking")

        logger.info(
            "Starting benchmark pipeline: %d model(s), %d run(s) each, judge=%s",
            len(models),
            self.runs,
            self.use_judge,
        )

        # Step 1: Elaborate intent
        logger.info("Elaborating intent: %r", intent)
        benchmark_intent = await self._elaborate_intent(intent)

        # Step 2: Run each model
        model_results: dict[str, list[ModelRunResult]] = {}
        for model in models:
            runs = []
            for run_idx in range(self.runs):
                logger.info("Running model %s (run %d/%d)", model, run_idx + 1, self.runs)
                result = await self.runner.run(
                    config_path, model, intent, model_params, central=self.main_config
                )
                runs.append(result)
            model_results[model] = runs

        # Step 3: Aggregate metrics
        logger.info("Aggregating metrics across all models")
        collector = MetricsCollector()
        benchmark_results: list[ModelBenchmarkResult] = []

        for model in models:
            runs = model_results[model]
            stats = collector.aggregate(model, runs)
            pricing = self.pricing.get_pricing(model)

            # Step 4: Compute cost
            estimated_cost = (
                stats.mean_prompt_tokens * pricing.input_cost_per_token
                + stats.mean_completion_tokens * pricing.output_cost_per_token
            )

            # Step 5: Judge quality (use first successful run's output)
            quality: QualityScore | None = None
            if self.use_judge and self.judge:
                successful_runs = [r for r in runs if r.success and r.output]
                if successful_runs:
                    logger.info("Judging output quality for model %s", model)
                    quality = await self.judge.evaluate(
                        intent=benchmark_intent,
                        output=successful_runs[0].output,
                        rubric=self.rubric,
                        model_name=model,
                    )

            benchmark_results.append(
                ModelBenchmarkResult(
                    model=model,
                    stats=stats,
                    quality=quality,
                    pricing=pricing,
                    estimated_cost=estimated_cost,
                )
            )

        # Step 6: Compute composite scores and rank
        self._compute_composite_scores(benchmark_results)
        logger.info("Benchmark complete: %d model(s) ranked", len(benchmark_results))

        return BenchmarkReport(
            intent=benchmark_intent,
            rubric=self.rubric,
            results=benchmark_results,
            evaluator_model=self.evaluator_model,
        )

    async def _elaborate_intent(self, intent: str) -> BenchmarkIntent:
        """Use evaluator model to elaborate the user's intent."""
        try:
            import json
            import re

            kwargs: dict[str, Any] = {
                "model": self.evaluator_model,
                "messages": [
                    {
                        "role": "user",
                        "content": INTENT_ELABORATION_PROMPT.format(intent=intent),
                    }
                ],
                "temperature": 0.0,
            }
            if litellm.supports_response_schema(self.evaluator_model, None):
                kwargs["response_format"] = {"type": "json_object"}

            response = await litellm.acompletion(**kwargs)

            raw = response.choices[0].message.content or "{}"
            # Strip markdown fences if model doesn't support json_object mode
            stripped = raw.strip()
            if stripped.startswith("```"):
                first_newline = stripped.find("\n")
                if first_newline != -1:
                    stripped = stripped[first_newline + 1 :]
                stripped = re.sub(r"```\s*$", "", stripped).strip()
                raw = stripped
            data = json.loads(raw)
            return BenchmarkIntent(
                raw_intent=intent,
                elaborated_intent=data.get("elaborated_intent", intent),
                expected_outcome=data.get("expected_outcome", ""),
                evaluation_criteria=data.get("evaluation_criteria", []),
            )
        except Exception as e:
            logger.warning("Intent elaboration failed, using raw intent: %s", e)
            return BenchmarkIntent(raw_intent=intent, elaborated_intent=intent)

    def _compute_composite_scores(self, results: list[ModelBenchmarkResult]) -> None:
        """Compute composite scores for ranking.

        Composite = quality_weight * normalized_quality
                   + cost_weight * normalized_cost_efficiency
                   + speed_weight * normalized_speed
        """
        if not results:
            return

        # Collect raw values
        quality_scores = [r.quality.overall_score if r.quality else 0.0 for r in results]
        costs = [r.estimated_cost for r in results]
        latencies = [r.stats.mean_latency_ms for r in results]

        # Normalize to 0-1 range (higher is better)
        norm_quality = _normalize_higher_better(quality_scores)
        norm_cost = _normalize_lower_better(costs)
        norm_speed = _normalize_lower_better(latencies)

        for i, result in enumerate(results):
            result.composite_score = round(
                self.quality_weight * norm_quality[i]
                + self.cost_weight * norm_cost[i]
                + self.speed_weight * norm_speed[i],
                4,
            )


def _normalize_higher_better(values: list[float]) -> list[float]:
    """Normalize values so higher is better, scaled to 0-1."""
    max_val = max(values) if values else 0
    if max_val == 0:
        return [0.0] * len(values)
    return [v / max_val for v in values]


def _normalize_lower_better(values: list[float]) -> list[float]:
    """Normalize values so lower original = higher normalized score, scaled to 0-1."""
    max_val = max(values) if values else 0
    if max_val == 0:
        return [1.0] * len(values)
    return [1.0 - (v / max_val) for v in values]
