"""
Localization Evaluator: Comprehensive checks for localized educational content.

Evaluates content against localization guidelines including:
- Schema fidelity (JSON structure preservation)
- Language correctness (grammar, spelling, terminology)
- Neutral scenario policy (cultural neutrality)
- Sensitivity guardrails (no religion, politics, etc.)
- Terminology & units (UAE curriculum standards, metric units)
- Tone & voice (appropriate educational tone)
- Clarity & cognitive load (minimal narrative overhead)
- Localization guardrail coverage
- Review checklist compliance
"""

from __future__ import annotations

import asyncio
import json
import logging
import textwrap
import time
from enum import Enum
from typing import Optional, Dict, Any

from pydantic import BaseModel

from .core.api_key_manager import get_openai_client

logger = logging.getLogger(__name__)

# Localization evaluation guidance based on InceptBench Localization Checklist
localization_guidance = textwrap.dedent("""
### Schema Fidelity
- **JSON keys remain in English**: All JSON keys (e.g., "question", "answer", "answer_options") must remain in English. No structural drift allowed.
- **Only value strings change**: Only the value strings should be translated/localized. Arrays and nesting structures must stay identical.
- **No structural changes**: The JSON schema structure must remain exactly the same as the original.

### Language Correctness
- **Target language compliance**: Content must be in the configured target language (e.g., Modern Standard Arabic for UAE).
- **Grammar and spelling**: Grammar, spelling, and punctuation must conform to grade-appropriate norms for the target language.
- **No leftover English**: There should be no leftover English text unless required (e.g., variables, units, acronyms like "AI", "DNA").
- **Grade-appropriate vocabulary**: Vocabulary should be appropriate for the target grade level.

### Neutral Scenario Policy
- **Generic contexts**: Contexts should be generic and universal: classroom, homework, simple shopping, basic measurements.
- **No cultural specifics**: No festivals, landmarks, public figures, or culturally specific props unless explicitly requested.
- **Universal solvability**: Problems should be solvable without requiring local cultural knowledge.
- **Neutral settings**: Avoid scenarios that assume specific cultural background knowledge.

### Sensitivity Guardrails
- **Zero sensitive content**: No references to religion, politics, dating, alcohol, gambling, or adult topics.
- **Gender balance**: Gender balance or gender-neutral phrasing where people appear in scenarios.
- **No stereotyping**: No stereotyping of emirates, regional groups, or any demographic groups.
- **Inclusive content**: Content should be inclusive and respectful of all backgrounds.

### Terminology & Units
- **UAE curriculum standards**: Educational terminology should follow UAE curriculum standards.
- **Metric units**: Units should be converted to metric system (meters, kilograms, liters, etc.).
- **Neutral currencies**: Currencies should be neutral unless instructions explicitly override (e.g., "AED" or generic "currency units").
- **Consistent date/time formats**: Date formats should be consistent (DD/MM/YYYY), time should use 24-hour clock format.
- **Appropriate terminology**: Subject-specific terminology should use standard terms for the target language/region.

### Tone & Voice
- **Friendly mentor tone**: Tone should be friendly, calm, and mentor-like with polite encouragement (e.g., "لنحاول معًا" - "Let's try together", "أحسنت" - "Well done").
- **No sarcasm**: No sarcasm or exaggerated enthusiasm.
- **Concise explanations**: Explanations should be concise and clear without forced cultural flavor.
- **Professional yet warm**: Maintain a professional educational tone that is warm and encouraging.

### Clarity & Cognitive Load
- **Minimal narrative context**: Students shouldn't need extra narrative context to understand the math/problem.
- **Short story text**: Story text in word problems should be short and focused.
- **Math-first approach**: Word problems should highlight the math rather than narrative flavor.
- **No ambiguous pronouns**: Avoid ambiguous pronouns or missing quantities that could confuse students.
- **Clear problem statement**: The problem should be immediately clear without requiring cultural or contextual knowledge.

### Localization Guardrail Coverage
- **Guardrail reflection**: All localization guardrails should be reflected in the output.
- **Neutral fallback**: If a guardrail can't be satisfied, the output should pick the neutral fallback instead of improvising.
- **Consistent application**: All guardrails should be consistently applied across all content.

### Review Checklist Compliance
- **Explicit verification**: Localization profile "QUALITY CHECKS" should be explicitly verified (e.g., "Ensure content uses only neutral, universally understandable scenarios").
- **Self-audit**: Self-audit should include schema compliance, unit consistency, and neutrality confirmation.
- **Comprehensive coverage**: All checklist items should be addressed.
""")

system_prompt = textwrap.dedent(f"""
You are an expert localization quality evaluator specializing in educational content localization, particularly for Arabic (UAE) and other target languages. Your task is to evaluate localized educational content against comprehensive localization guidelines.

## Evaluation Framework

Evaluate the content against the following localization criteria:

{localization_guidance}

## Evaluation Process

For each criterion, provide:
1. **Pass/Fail status**: PASS if the content fully meets the criterion, FAIL if it does not
2. **Score**: A score from 0-2 where:
   - 0 = FAIL (criterion not met)
   - 1 = PARTIAL (mostly meets but has minor issues)
   - 2 = PASS (fully meets criterion)
3. **Reasoning**: Brief explanation of why the criterion passed or failed
4. **Issues**: Specific issues found (if any)
5. **Strengths**: What was done well (if applicable)

## Overall Assessment

Provide an overall localization quality score (0-1) based on:
- Weighted average of all criteria scores (normalized to 0-1)
- Critical failures (schema fidelity, sensitivity guardrails) should heavily penalize the score
- Language correctness and neutrality are high-priority criteria

## Output Format

Return a JSON object with:
- Individual criterion scores (0-2) and pass/fail status
- Overall score (0-1)
- List of issues found
- List of strengths
- Risk notes (free text) for any flagged concerns
- Recommendation: "accept", "revise", or "reject"

## Instructions

The user will provide:
1. The content to evaluate (question, text content, or article)
2. The target language (e.g., "ar" for Arabic)
3. The original content (if available) for schema comparison
4. Any additional context (grade, subject, etc.)

Evaluate the content comprehensively and provide detailed feedback.
""")

class LocalizationCriterionResult(BaseModel):
    """Result for a single localization criterion"""
    criterion: str
    pass_fail: str  # "PASS" or "FAIL"
    score: int  # 0, 1, or 2
    reasoning: str
    issues: list[str]
    strengths: list[str]

class LocalizationEvaluationResult(BaseModel):
    """Complete localization evaluation result"""
    # Schema Fidelity
    schema_fidelity: LocalizationCriterionResult
    
    # Language Correctness
    language_correctness: LocalizationCriterionResult
    
    # Neutral Scenario Policy
    neutral_scenario: LocalizationCriterionResult
    
    # Sensitivity Guardrails
    sensitivity_guardrails: LocalizationCriterionResult
    
    # Terminology & Units
    terminology_units: LocalizationCriterionResult
    
    # Tone & Voice
    tone_voice: LocalizationCriterionResult
    
    # Clarity & Cognitive Load
    clarity_cognitive_load: LocalizationCriterionResult
    
    # Localization Guardrail Coverage
    guardrail_coverage: LocalizationCriterionResult
    
    # Review Checklist Compliance
    checklist_compliance: LocalizationCriterionResult
    
    # Overall assessment
    overall_score: float  # 0-1
    recommendation: str  # "accept", "revise", or "reject"
    issues: list[str]
    strengths: list[str]
    risk_notes: str  # Free text for flagged concerns

async def evaluate_localization(
    content: str,
    target_language: str = "ar",
    target_locale: Optional[str] = None,
    cultural_context: Optional[str] = None,
    content_type: str = "question",
    original_content: Optional[str] = None,
    grade: Optional[str] = None,
    subject: Optional[str] = None,
    skill_title: Optional[str] = None
) -> str:
    """
    Evaluates localized educational content using OpenAI's o3 model.
    
    Parameters
    ----------
    content : str
        The localized content to evaluate (question text, answer, explanation, etc.)
    target_language : str
        Target language code (e.g., "ar" for Arabic, "en" for English)
    target_locale : Optional[str]
        Locale identifier (e.g., "ar-AE" for Arabic/UAE, "en-IN" for English/India)
    cultural_context : Optional[str]
        Cultural region (e.g., "AE", "IN") to guide cultural guardrails
    content_type : str
        Type of content: "question", "text", or "article"
    original_content : Optional[str]
        Original content for schema comparison (if available)
    grade : Optional[str]
        Target grade level
    subject : Optional[str]
        Subject area
    skill_title : Optional[str]
        Skill/topic title
        
    Returns
    -------
    str
        The evaluation response from the model as a JSON string
    """
    logger.info(
        f"Evaluating localization for {content_type} in {target_language}"
        f"{f' ({target_locale})' if target_locale else ''}: {content[:50]}..."
    )
    
    # Build evaluation prompt
    user_prompt = f"""Evaluate the following localized educational content:

**Content Type**: {content_type}
**Target Language**: {target_language}
**Target Locale**: {target_locale or "Not specified"}
**Cultural Context**: {cultural_context or "Not specified"}
**Grade**: {grade or "Not specified"}
**Subject**: {subject or "Not specified"}
**Topic**: {skill_title or "Not specified"}

**Content to Evaluate**:
{content}
"""

    if original_content:
        user_prompt += f"""

**Original Content (for schema comparison)**:
{original_content}
"""

    user_prompt += """

Please evaluate this content against all localization criteria and provide a comprehensive assessment.
"""
    
    def _sync_evaluate_with_retry():
        """Synchronous function to call OpenAI API with retry logic"""
        max_retries = 2
        
        for attempt in range(max_retries + 1):
            try:
                client = get_openai_client(timeout=300.0)  # 5 minutes timeout
                
                # Use the responses.parse API for structured output with o3
                response = client.responses.parse(
                    model="o3",
                    input=[
                        {"role": "developer", "content": system_prompt},
                        {"role": "user", "content": user_prompt}
                    ],
                    text_format=LocalizationEvaluationResult
                )
                return response
                
            except Exception as e:
                attempt_num = attempt + 1
                logger.warning(
                    f"Error in localization evaluation (attempt {attempt_num}/{max_retries + 1}): "
                    f"{type(e).__name__}: {str(e)}"
                )
                
                if attempt < max_retries:
                    backoff_time = min(2 ** attempt, 10)  # Exponential backoff, max 10 seconds
                    logger.info(
                        f"Localization evaluation waiting {backoff_time}s before retry {attempt_num + 1}"
                    )
                    time.sleep(backoff_time)
                else:
                    logger.error(f"Localization evaluation failed after {max_retries + 1} attempts")
                    raise
        
        return None
    
    try:
        # Run the OpenAI call in a thread pool to avoid blocking the event loop
        response = await asyncio.get_event_loop().run_in_executor(None, _sync_evaluate_with_retry)
        
        if response is None:
            return json.dumps({"error": "Evaluation failed after retries"})
        
        # Process the structured response
        for output_item in response.output:
            if output_item.type == "message":
                # Look for parsed data in the message content
                for content_item in output_item.content:
                    if (content_item.type == "output_text" and 
                        hasattr(content_item, "parsed") and 
                        content_item.parsed is not None):
                        
                        if hasattr(content_item, "text") and content_item.text:
                            try:
                                return content_item.text
                            except Exception as text_error:
                                logger.warning(
                                    f"Error using text representation: {str(text_error)}"
                                )
                        
                        # Manual conversion of the parsed object to handle enums
                        try:
                            # Convert parsed object to dict
                            parsed_dict = content_item.parsed.model_dump()
                            return json.dumps(parsed_dict)
                        except Exception as enum_error:
                            logger.warning(f"Error handling enum serialization: {str(enum_error)}")
                
                # Fallback to the raw text if parsed object not available but text is
                for content_item in output_item.content:
                    if content_item.type == "output_text" and hasattr(content_item, "text"):
                        try:
                            return content_item.text
                        except Exception as text_error:
                            logger.warning(f"Error using text: {str(text_error)}")
                            return json.dumps({"error": "Could not parse structured output"})
        
        # If we can't find anything, return an error
        logger.error("Could not extract localization evaluation from response")
        return json.dumps({"error": "Could not extract evaluation results"})
        
    except Exception as e:
        logger.error(f"Error evaluating localization: {str(e)}")
        return json.dumps({"error": f"Localization evaluation failed: {str(e)}"})

