"""
LLM-based grader for evaluating agent rationales and reasoning quality.
Uses OpenAI-compatible API to grade soft criteria.
"""

import os
import json
import logging
from typing import Any
from pydantic import BaseModel, Field

from openai import AsyncOpenAI

from env.state import EnvState

logger = logging.getLogger(__name__)


class LLMGradeResult(BaseModel):
    """Result from LLM grading."""
    diagnostic_coherence: float = Field(ge=0.0, le=1.0)
    evidence_sufficiency: float = Field(ge=0.0, le=1.0)
    root_cause_narration: float = Field(ge=0.0, le=1.0)
    overall: float = Field(ge=0.0, le=1.0)
    reasoning: str
    
    model_config = {"frozen": False}


class LLMGrader:
    """
    LLM judge using OpenAI-compatible API.
    Evaluates agent reasoning and rationales.
    """
    
    SYSTEM_PROMPT = """You are an expert IT operations grader evaluating an AI agent's performance on an incident triage task.

You will be given:
1. The hidden root cause (ground truth)
2. The agent's action history with observation updates
3. The final outcome (resolved/escalated/closed)

Grade the following dimensions on a scale of 0.0 to 1.0:
- diagnostic_coherence: Are the agent's actions logically consistent with the observations? Did it follow a coherent diagnostic strategy?
- evidence_sufficiency: Did the agent gather enough evidence before committing to a fix or escalation?
- root_cause_narration: Did the agent's actions demonstrate understanding of the actual root cause?

Return ONLY a valid JSON object with these exact keys:
{
  "diagnostic_coherence": <float 0.0-1.0>,
  "evidence_sufficiency": <float 0.0-1.0>,
  "root_cause_narration": <float 0.0-1.0>,
  "overall": <float 0.0-1.0>,
  "reasoning": "<brief explanation>"
}

Do not include any other text or markdown formatting."""
    
    def __init__(self):
        """Initialize LLM grader with API configuration."""
        self.api_base_url = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
        self.model_name = os.environ.get("MODEL_NAME", "gpt-4o-mini")
        self.api_key = os.environ.get("OPENAI_API_KEY", "")
        
        if not self.api_key:
            logger.warning("OPENAI_API_KEY not set. LLM grading will use fallback scores.")
        
        self.client = AsyncOpenAI(
            api_key=self.api_key,
            base_url=self.api_base_url
        )
    
    async def grade(self, state: EnvState) -> LLMGradeResult:
        """
        Grade episode using LLM judge.
        
        Args:
            state: Complete environment state after episode
        
        Returns:
            LLMGradeResult with scores and reasoning
        """
        if not self.api_key:
            logger.warning("No API key available, returning fallback scores")
            return self._fallback_grade()
        
        try:
            # Build user prompt
            user_prompt = self._build_prompt(state)
            
            # Call LLM API
            response = await self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": self.SYSTEM_PROMPT},
                    {"role": "user", "content": user_prompt}
                ],
                temperature=0.3,
                max_tokens=500
            )
            
            # Parse response
            content = response.choices[0].message.content
            if not content:
                logger.error("Empty response from LLM")
                return self._fallback_grade()
            
            # Try to parse JSON
            try:
                result_dict = json.loads(content)
            except json.JSONDecodeError:
                # Try to extract JSON from markdown code blocks
                if "```json" in content:
                    content = content.split("```json")[1].split("```")[0].strip()
                elif "```" in content:
                    content = content.split("```")[1].split("```")[0].strip()
                result_dict = json.loads(content)
            
            return LLMGradeResult(**result_dict)
        
        except Exception as e:
            logger.error(f"LLM grading failed: {e}", exc_info=True)
            return self._fallback_grade()
    
    def _build_prompt(self, state: EnvState) -> str:
        """Build grading prompt from episode state."""
        # Extract key information
        root_cause = state.hidden.root_cause
        category = state.hidden.root_cause_category
        severity = state.hidden.severity
        
        # Summarize action sequence
        action_summary = []
        for i, action in enumerate(state.action_log, 1):
            action_summary.append(
                f"{i}. {action.action_name} (reward: {action.reward_received:.2f})"
            )
        
        actions_text = "\n".join(action_summary)
        
        # Summarize observations
        final_obs = state.observation
        
        prompt = f"""## Episode Summary

**Ground Truth Root Cause:**
- Category: {category}
- Description: {root_cause}
- Severity: {severity}
- Affected Users: {state.hidden.affected_users_count}

**Correct Remediation:**
- Primary: {state.hidden.correct_remediation}
- Alternatives: {', '.join(state.hidden.correct_remediation_alts) if state.hidden.correct_remediation_alts else 'None'}

**Agent's Action Sequence:**
{actions_text}

**Final State:**
- Steps taken: {final_obs.step_count}
- Episode terminated: {state.done}
- Diagnostic steps: {state.diagnostic_steps_taken}
- Escalation allowed: {final_obs.escalation_allowed}

**Observations:**
- Network Status: {final_obs.network_status.value}
- VPN Status: {final_obs.vpn_status.value}
- Auth Status: {final_obs.auth_status.value}
- Service Health: {dict(final_obs.service_health)}

**Performance:**
- Cumulative Reward: {state.cumulative_reward:.2f}
- Task Level: {state.task_level.value}

Evaluate how well the agent diagnosed and resolved this incident."""
        
        return prompt
    
    def _fallback_grade(self) -> LLMGradeResult:
        """Return neutral fallback scores when LLM unavailable."""
        return LLMGradeResult(
            diagnostic_coherence=0.5,
            evidence_sufficiency=0.5,
            root_cause_narration=0.5,
            overall=0.5,
            reasoning="LLM grading unavailable, using fallback scores"
        )