"""LLM-as-judge for evaluating agent outputs against rubrics."""

from __future__ import annotations

import json
import logging
import re
from typing import Any

import litellm

from sage_evaluator.exceptions import EvaluationError
from sage_evaluator.models import (
    BenchmarkIntent,
    DimensionScore,
    EvaluationRubric,
    QualityScore,
)

logger = logging.getLogger(__name__)

JUDGE_SYSTEM_PROMPT = """\
You are an expert evaluator assessing the quality of an AI agent's output.
You will be given:
1. The user's intent (what was asked)
2. Expected outcome criteria
3. The agent's actual output
4. A rubric with scoring dimensions

Score each dimension from 1 to 5:
- 1: Very poor, fails to meet the criteria
- 2: Below average, partially meets criteria with significant gaps
- 3: Average, meets basic criteria but has notable shortcomings
- 4: Good, meets criteria well with minor issues
- 5: Excellent, fully meets or exceeds criteria

You MUST respond with valid JSON matching this exact structure:
{
  "dimension_scores": [
    {
      "dimension": "<dimension name>",
      "score": <1-5>,
      "reasoning": "<brief explanation>"
    }
  ]
}

Be objective and critical. Do not inflate scores.
Evaluate based solely on the rubric dimensions provided.
"""


def _strip_markdown_fences(text: str) -> str:
    """Remove markdown code fences (```json ... ```) if present."""
    stripped = text.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()
    return stripped


class LLMJudge:
    """Evaluates agent outputs using an LLM as judge with structured rubric scoring."""

    def __init__(self, model: str):
        self.model = model

    async def evaluate(
        self,
        intent: BenchmarkIntent,
        output: str,
        rubric: EvaluationRubric,
        model_name: str = "",
    ) -> QualityScore:
        """Evaluate an agent's output against a rubric.

        Args:
            intent: The elaborated benchmark intent with criteria.
            output: The agent's output to evaluate.
            rubric: The rubric to score against.
            model_name: Name of the model that produced the output.

        Returns:
            QualityScore with dimension scores and weighted overall score.
        """
        logger.info("Evaluating output for model %r against rubric %r", model_name, rubric.name)
        user_prompt = self._build_evaluation_prompt(intent, output, rubric)

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

        try:
            response = await litellm.acompletion(**kwargs)
        except Exception as e:
            raise EvaluationError(f"Judge LLM call failed: {e}") from e

        raw_output = response.choices[0].message.content or ""
        raw_output = _strip_markdown_fences(raw_output)
        dimension_scores = self._parse_scores(raw_output, rubric)
        overall_score = self._compute_weighted_score(dimension_scores, rubric)
        logger.info(
            "Judge evaluation complete for model %r: overall score %.2f/5",
            model_name,
            overall_score,
        )

        return QualityScore(
            model=model_name,
            overall_score=overall_score,
            dimension_scores=dimension_scores,
            raw_judge_output=raw_output,
        )

    def _build_evaluation_prompt(
        self,
        intent: BenchmarkIntent,
        output: str,
        rubric: EvaluationRubric,
    ) -> str:
        """Build the evaluation prompt for the judge."""
        criteria_text = "\n".join(f"- {c}" for c in intent.evaluation_criteria) or "None specified"

        dimensions_text = "\n".join(
            f"- **{d.name}** (weight: {d.weight}): {d.description}" for d in rubric.dimensions
        )

        return f"""\
## User Intent
{intent.raw_intent}

## Elaborated Intent
{intent.elaborated_intent}

## Expected Outcome
{intent.expected_outcome}

## Evaluation Criteria
{criteria_text}

## Agent Output
```
{output}
```

## Rubric: {rubric.name}
{rubric.description}

### Dimensions
{dimensions_text}

Score each dimension from 1-5 with reasoning. Respond with JSON only.
"""

    def _parse_scores(self, raw_output: str, rubric: EvaluationRubric) -> list[DimensionScore]:
        """Parse dimension scores from judge JSON output."""
        try:
            data = json.loads(raw_output)
        except json.JSONDecodeError as e:
            logger.warning("Failed to parse judge output as JSON: %s", e)
            return [
                DimensionScore(dimension=d.name, score=0.0, reasoning="Parse error")
                for d in rubric.dimensions
            ]

        raw_scores = data.get("dimension_scores", [])

        # Some models return dimension_scores as a JSON string instead of a list
        if isinstance(raw_scores, str):
            logger.warning("dimension_scores is a string, attempting to parse as JSON")
            try:
                raw_scores = json.loads(raw_scores)
            except json.JSONDecodeError:
                logger.warning("Failed to parse dimension_scores string as JSON")
                raw_scores = []

        if not isinstance(raw_scores, list):
            logger.warning("dimension_scores is not a list: %s", type(raw_scores).__name__)
            raw_scores = []

        score_map: dict[str, DimensionScore] = {}

        for item in raw_scores:
            if not isinstance(item, dict):
                logger.warning("Skipping non-dict score item: %r", item)
                continue
            name = item.get("dimension", "")
            try:
                score = float(item.get("score", 0))
            except (TypeError, ValueError):
                score = 0.0
            score = max(1.0, min(5.0, score))  # clamp to 1-5
            reasoning = item.get("reasoning", "")
            score_map[name] = DimensionScore(dimension=name, score=score, reasoning=reasoning)

        # Ensure all rubric dimensions are present
        result = []
        for dim in rubric.dimensions:
            if dim.name in score_map:
                result.append(score_map[dim.name])
            else:
                result.append(
                    DimensionScore(dimension=dim.name, score=0.0, reasoning="Not scored by judge")
                )
        return result

    def _compute_weighted_score(
        self,
        dimension_scores: list[DimensionScore],
        rubric: EvaluationRubric,
    ) -> float:
        """Compute weighted overall score from dimension scores.

        The overall score is NOT from the judge's self-reported number.
        It is computed client-side as a weighted average.
        """
        weight_map = {d.name: d.weight for d in rubric.dimensions}
        total_weight = sum(weight_map.values())

        if total_weight == 0:
            return 0.0

        weighted_sum = sum(ds.score * weight_map.get(ds.dimension, 1.0) for ds in dimension_scores)
        return round(weighted_sum / total_weight, 2)
