"""
LLM-based judge using LiteLLM
"""

from openevalkit.judges.base import Judge
from openevalkit.run import Run
from openevalkit.judges.rubric import Rubric
from openevalkit.judgment import Judgment
from openevalkit.judges.llm_config import LLMConfig
from openevalkit.judges.prompt import PromptTemplate, DEFAULT_JUDGE_TEMPLATE
from openevalkit.errors import JudgingError
import litellm
from typing import Optional
import json
import re

class LLMJudge(Judge):
    """
    LLM-based judge using LiteLLM for provider-agnostic LLM calls.
    
    Supports 100+ models: OpenAI, Anthropic, Google, Cohere, Ollama (local), and more.
    
    **Note on Confidence Scores:**
    LLMs often struggle with accurate self-assessment of confidence. The 
    confidence values returned may not be well-calibrated and should be 
    interpreted with caution. Research shows:
    
    - LLMs tend to be overconfident (report high confidence even when uncertain)
    - Confidence often correlates highly with the score itself
    - Without specialized calibration training, confidence is less reliable
    
    **When confidence might be useful:**
    - Filtering low-confidence judgments for manual review
    - Comparing confidence across different models/prompts
    - Research on LLM calibration and uncertainty
    
    **When to ignore confidence:**
    - Making critical decisions (use multiple judges instead)
    - Comparing scores across different criteria
    - When you need well-calibrated probabilities
    
    For more reliable uncertainty estimates, consider using ensemble judges
    with multiple models and examining score variance.

    Examples:
        >>> from openevalkit.judges import LLMConfig, Rubric, LLMJudge
        >>> 
        >>> # OpenAI
        >>> config = LLMConfig(model="gpt-4o")
        >>> rubric = Rubric(criteria=["helpfulness", "clarity"], scale="0-1")
        >>> judge = LLMJudge(llm_config=config, rubric=rubric)
        >>> 
        >>> # Local Ollama
        >>> config = LLMConfig(model="ollama/llama3", api_base="http://localhost:11434")
        >>> judge = LLMJudge(llm_config=config, rubric=rubric)
        >>>
        >>> # Basic usage (confidence included by default)
        >>> judge = LLMJudge(
        ...     llm_config=LLMConfig(model="gpt-4o"),
        ...     rubric=Rubric(criteria=["helpfulness"], scale="0-1")
        ... )
        >>> judgment = judge.judge(run, rubric)
        >>> judgment.confidence  # Available but use cautiously
        0.85
    """

    def __init__(
            self,
            llm_config: LLMConfig,
            rubric: Rubric,
            prompt_template: Optional[PromptTemplate] = None,
            confidence_method: str = "weighted_avg"
            ):
        """
        Initialize LLM judge.

        Args:
            llm_config: LLM configuration (model, API settings, temperature, etc.)
            rubric: Evaluation criteria with optional weights
            prompt_template: Custom prompt (uses DEFAULT_JUDGE_TEMPLATE if None)
            confidence_method: How to aggregate confidence ("weighted_avg", "min", "variance_based")
        """

        self.llm_config = llm_config
        self.rubric = rubric
        self.prompt_template = prompt_template or DEFAULT_JUDGE_TEMPLATE
        self.confidence_method = confidence_method
        self.name = f"llm_judge_{self.llm_config.model.replace('/', '_')}"


    def judge(self, run: Run, rubric: Rubric) -> Judgment:
        """
        Judge a single run using LLM.
        
        Args:
            run: Run to judge
            rubric: Evaluation criteria (can override instance rubric)
            
        Returns:
            Judgment with scores, explanations, and confidence
            
        Raises:
            JudgingError: If judging fails
        """

        try:
            # Build the prompt
            prompt = self._build_prompt(run, rubric)

            # call LLM
            response = litellm.completion(
                model=self.llm_config.model,
                messages=[{'role': 'user', 'content': prompt}],
                temperature=self.llm_config.temperature,
                max_tokens=self.llm_config.max_tokens,
                api_key=self.llm_config.api_key,
                api_base=self.llm_config.api_base,
                num_retries=self.llm_config.num_retries,
                retry_after=self.llm_config.retry_after,
                **self.llm_config.extra_kwargs
            )

            # Parse response into judgement
            judgment = self._parse_response_into_judgment(response, rubric)

            return judgment
        
        except Exception as e:
            raise JudgingError(
                f"Failed to judge run {run.id} with {self.llm_config.model}: {str(e)}"
                ) 

    
    def _build_prompt(self, run: Run, rubric: Rubric) -> str:
        """Build evaluation prompt from template."""

        return self.prompt_template.render(
            input=str(run.input),
            output=run.output,
            reference=str(run.reference) if run.reference else "None",
            rubric=rubric.to_prompt() 
            )
    
    def _parse_response_into_judgment(self, response, rubric: Rubric) -> Judgment:
        """
        Parse LLM output into a judgment object
        """
        # Extract content
        content = response.choices[0].message.content

        # Parse JSON (with error for markdown code blocks)
        try:
            data = json.loads(content) 

        except json.JSONDecodeError:
            # Try to extract JSON from markdown code blocks
            json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', content, re.DOTALL)
            if json_match:
                data = json.loads(json_match.group(1))
            else:
                raise JudgingError(f"Failed to parse LLM response as JSON: {content}")
        
        # Extract per criteria data
        criteria_scores = data.get("criteria_scores", {})
        criteria_explanations = data.get("criteria_explanations", {})
        criteria_confidences = data.get("criteria_confidences")
        
        # Get overall explanation
        explanation = data.get("explanation", "")

        # Compute overall score (framework computes this, not LLM)
        overall_score = None
        if criteria_scores:
            overall_score = self._compute_overall_score(criteria_scores, rubric.weights)

        # Compute overall confidence (framework computes this, not LLM)
        overall_confidence = None
        if criteria_confidences:
            overall_confidence = self._compute_overall_confidence(criteria_confidences=criteria_confidences, weights=rubric.weights, method=self.confidence_method)
        
        return Judgment(
            score=overall_score,
            criteria_scores=criteria_scores,
            criteria_explanations=criteria_explanations,
            explanation=explanation,
            criteria_confidences=criteria_confidences,
            confidence=overall_confidence,
            metadata={
            "model": self.llm_config.model,
            "temperature": self.llm_config.temperature,
            "prompt_hash": self.prompt_template.hash(),
            "raw_response": content
            }

        )
        
