"""LLM-as-judge scoring for open-ended evaluation.

This module provides LLM-based quality assessment for responses that can't
be evaluated with exact match or F1. Uses a judge model (e.g., claude-opus-4.5)
to score responses on quality dimensions.
"""

import asyncio
import logging
import re
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)

# Default judge prompts for different evaluation types
JUDGE_PROMPTS = {
    "quality": """You are an expert evaluator. Score the following response on a scale of 1-5.

Question: {question}

Response to evaluate: {response}

{reference_section}

Scoring criteria:
1 = Completely incorrect, irrelevant, or harmful
2 = Mostly incorrect with some relevant information
3 = Partially correct, addresses the question but with errors
4 = Mostly correct with minor issues
5 = Excellent response, accurate and complete

Provide your score as a single number (1-5) followed by a brief justification.
Format: SCORE: [number]
JUSTIFICATION: [your reasoning]""",
    "comparison": """You are an expert evaluator comparing two responses.

Question: {question}

Response A: {response_a}

Response B: {response_b}

Which response is better? Consider accuracy, completeness, and clarity.
Options: A, B, or TIE

Format: WINNER: [A/B/TIE]
JUSTIFICATION: [your reasoning]""",
    "factuality": """You are a fact-checker. Evaluate if the response is factually accurate.

Question: {question}

Response: {response}

Reference answer (ground truth): {reference}

Is the response factually correct? Consider:
- Core facts are accurate
- No contradictions with the reference
- No hallucinated information

Format: VERDICT: [CORRECT/INCORRECT/PARTIALLY_CORRECT]
SCORE: [0-1, where 1 is fully correct]
JUSTIFICATION: [your reasoning]""",
    "helpfulness": """You are an expert evaluator assessing helpfulness.

User request: {question}

Assistant response: {response}

Score how helpful and useful this response is (1-5):
1 = Not helpful at all, doesn't address the request
2 = Minimally helpful, significant gaps
3 = Moderately helpful, addresses main points
4 = Very helpful, comprehensive with minor gaps
5 = Extremely helpful, exceeds expectations

Format: SCORE: [1-5]
JUSTIFICATION: [your reasoning]""",
}


@dataclass
class JudgeResult:
    """Result from LLM judge evaluation.

    Attributes:
        score: Numeric score (0-1 normalized, or raw 1-5).
        normalized_score: Score normalized to 0-1 range.
        verdict: Categorical verdict if applicable.
        justification: Judge's reasoning.
        raw_response: Full judge response.
        judge_model: Model used for judging.
        metadata: Additional evaluation metadata.
    """

    score: float
    normalized_score: float
    verdict: str | None = None
    justification: str | None = None
    raw_response: str = ""
    judge_model: str = ""
    metadata: dict[str, Any] = field(default_factory=dict)


class LLMJudge:
    """LLM-based judge for evaluating response quality.

    Uses a capable LLM (e.g., claude-opus-4.5, gpt-5.1) to assess the quality
    of model outputs. Essential for benchmarks like Arena-Hard where quality
    isn't easily quantified.

    Example:
        judge = LLMJudge(model="claude-opus-4.5")
        result = await judge.score(
            question="Explain quantum entanglement",
            response="Quantum entanglement is when two particles...",
            reference="[optional reference answer]",
        )
        print(f"Score: {result.normalized_score:.2f}")
    """

    def __init__(
        self,
        model: str = "claude-opus-4.5",
        prompt_type: str = "quality",
        custom_prompt: str | None = None,
        max_retries: int = 3,
        temperature: float = 0.0,
    ):
        """Initialize LLM judge.

        Args:
            model: Judge model to use.
            prompt_type: Type of prompt ("quality", "comparison", "factuality").
            custom_prompt: Optional custom prompt template.
            max_retries: Maximum retry attempts on failure.
            temperature: Sampling temperature (0 for deterministic).
        """
        self._model = model
        self._prompt_type = prompt_type
        self._prompt_template = custom_prompt or JUDGE_PROMPTS.get(
            prompt_type, JUDGE_PROMPTS["quality"]
        )
        self._max_retries = max_retries
        self._temperature = temperature

    def _format_prompt(
        self,
        question: str,
        response: str,
        reference: str | None = None,
        response_b: str | None = None,
    ) -> str:
        """Format the judge prompt with provided values."""
        reference_section = ""
        if reference:
            reference_section = f"Reference answer: {reference}"

        return self._prompt_template.format(
            question=question,
            response=response,
            response_a=response,
            response_b=response_b or "",
            reference=reference or "",
            reference_section=reference_section,
        )

    def _parse_score(self, response: str) -> tuple[float, str | None, str | None]:
        """Parse score, verdict, and justification from judge response.

        Returns:
            Tuple of (score, verdict, justification).
        """
        response = response.strip()

        # Extract score
        score_match = re.search(r"SCORE:\s*(\d+(?:\.\d+)?)", response, re.IGNORECASE)
        score = float(score_match.group(1)) if score_match else 0.0

        # Extract verdict
        verdict_match = re.search(
            r"(?:VERDICT|WINNER):\s*([A-Z_]+)",
            response,
            re.IGNORECASE,
        )
        verdict = verdict_match.group(1).upper() if verdict_match else None

        # Extract justification
        just_match = re.search(
            r"JUSTIFICATION:\s*(.+?)(?:\n|$)",
            response,
            re.IGNORECASE | re.DOTALL,
        )
        justification = just_match.group(1).strip() if just_match else None

        return score, verdict, justification

    async def score(
        self,
        question: str,
        response: str,
        reference: str | None = None,
        **kwargs: Any,
    ) -> JudgeResult:
        """Score a single response using the LLM judge.

        Args:
            question: Original question/prompt.
            response: Response to evaluate.
            reference: Optional reference answer.
            **kwargs: Additional prompt format arguments.

        Returns:
            JudgeResult with score and justification.
        """
        from ember.api.models import models as models_api

        prompt = self._format_prompt(question, response, reference)

        for attempt in range(self._max_retries):
            try:
                judge_response = await asyncio.to_thread(
                    models_api.response,
                    self._model,
                    prompt,
                    temperature=self._temperature,
                )

                # models_api.response() returns Response with .text
                raw = judge_response.text

                score, verdict, justification = self._parse_score(raw)

                # Normalize score to 0-1
                if self._prompt_type in ("quality", "helpfulness"):
                    normalized = score / 5.0  # 1-5 scale
                elif self._prompt_type == "factuality":
                    normalized = score  # Already 0-1
                else:
                    normalized = score / 5.0

                return JudgeResult(
                    score=score,
                    normalized_score=min(1.0, max(0.0, normalized)),
                    verdict=verdict,
                    justification=justification,
                    raw_response=raw,
                    judge_model=self._model,
                    metadata={"attempt": attempt + 1},
                )

            except Exception as e:
                logger.warning(
                    f"Judge attempt {attempt + 1} failed: {e}",
                    exc_info=True,
                )
                if attempt == self._max_retries - 1:
                    return JudgeResult(
                        score=0.0,
                        normalized_score=0.0,
                        verdict="ERROR",
                        justification=str(e),
                        judge_model=self._model,
                        metadata={"error": str(e)},
                    )
                await asyncio.sleep(2 ** attempt)  # Exponential backoff

        # Should not reach here
        return JudgeResult(score=0.0, normalized_score=0.0, judge_model=self._model)

    async def compare(
        self,
        question: str,
        response_a: str,
        response_b: str,
    ) -> JudgeResult:
        """Compare two responses and determine winner.

        Args:
            question: Original question/prompt.
            response_a: First response.
            response_b: Second response.

        Returns:
            JudgeResult with verdict (A, B, or TIE).
        """
        from ember.api.models import models as models_api

        prompt = self._format_prompt(
            question,
            response_a,
            response_b=response_b,
        )

        try:
            judge_response = await asyncio.to_thread(
                models_api.response,
                self._model,
                prompt,
                temperature=self._temperature,
            )

            # models_api.response() returns Response with .text
            raw = judge_response.text

            _, verdict, justification = self._parse_score(raw)

            # Convert verdict to score
            score = 0.5  # Default tie
            if verdict == "A":
                score = 1.0
            elif verdict == "B":
                score = 0.0

            return JudgeResult(
                score=score,
                normalized_score=score,
                verdict=verdict,
                justification=justification,
                raw_response=raw,
                judge_model=self._model,
            )

        except Exception as e:
            logger.error(f"Comparison judge failed: {e}", exc_info=True)
            return JudgeResult(
                score=0.5,
                normalized_score=0.5,
                verdict="ERROR",
                justification=str(e),
                judge_model=self._model,
                metadata={"error": str(e)},
            )

    async def score_batch(
        self,
        examples: list[dict[str, str]],
        max_concurrent: int = 5,
    ) -> list[JudgeResult]:
        """Score multiple examples with controlled concurrency.

        Args:
            examples: List of dicts with "question", "response", "reference".
            max_concurrent: Maximum concurrent judge calls.

        Returns:
            List of JudgeResult, one per example.
        """
        semaphore = asyncio.Semaphore(max_concurrent)

        async def score_one(example: dict[str, str]) -> JudgeResult:
            async with semaphore:
                return await self.score(
                    question=example["question"],
                    response=example["response"],
                    reference=example.get("reference"),
                )

        tasks = [score_one(ex) for ex in examples]
        return await asyncio.gather(*tasks)


__all__ = [
    "JudgeResult",
    "LLMJudge",
    "JUDGE_PROMPTS",
]
