"""AI-powered quality calibration judge for robot training datasets.

Uses Gemini 2.5 Flash as a binary calibration judge to verify whether
A-graded datasets truly deserve the top tier. Research-backed design:

- Binary verification > numerical scoring (more reliable)
- Few-shot anchor examples for calibration (+25-30% accuracy)
- Chain-of-thought reasoning before verdict
- Selective invocation: only called for A-graded datasets (~$0.001/call)

References:
    - "Trust or Escalate" (ICLR 2025): selective evaluation with confidence
    - Belkhale et al. (NeurIPS 2023): action divergence as primary quality axis
    - LLM-as-Judge best practices: binary + CoT + few-shot anchors
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass

logger = logging.getLogger(__name__)


@dataclass
class AIJudgeResult:
    """Result of AI quality calibration check."""

    keep_grade: bool  # True = keep current grade, False = downgrade
    reasoning: str  # Chain-of-thought explanation
    model_used: str
    error: str | None = None


# The calibrated binary judge prompt. Key design decisions:
#   1. Binary question (keep A or downgrade to B), not open-ended scoring
#   2. Three anchor examples calibrated to prevent systematic harshness
#   3. Reasoning requested BEFORE verdict
#   4. Explicit instruction that A is correct for MOST clean datasets
_BINARY_JUDGE_PROMPT = (
    "You are an expert robotics data quality judge. I will show you statistical "
    "diagnostics for a robot training dataset that our automated system graded "
    'as "A" (top quality). Your ONE task: decide if "A" is appropriate, or if '
    '"B" (good but minor issues) fits better.\n\n'
    "CALIBRATION EXAMPLES:\n\n"
    "Example 1 — A IS appropriate:\n"
    "  Divergence: 0.10 (low), Clipping: 0, PolicyFit: 0.90, ElevatedCVs: 1/8\n"
    "  → VERDICT: keep A. Low divergence, minimal issues.\n\n"
    "Example 2 — B is more appropriate:\n"
    "  Divergence: 0.16 (moderate), Clipping: 2, PolicyFit: 0.88, ElevatedCVs: 2/8\n"
    "  → VERDICT: downgrade to B. Moderate divergence plus clipping = compounding "
    "concerns. Demonstrations vary in approach.\n\n"
    "Example 3 — A IS appropriate despite some signals:\n"
    "  Divergence: 0.06 (very low), Clipping: 0, PolicyFit: 0.77, ElevatedCVs: 3/8\n"
    "  → VERDICT: keep A. Very low divergence means data is clean. Elevated CVs "
    "reflect intentional trajectory variation, not quality problems.\n\n"
    "IMPORTANT: A is correct for most clean datasets. Only downgrade when you see "
    "SPECIFIC compounding evidence (divergence >= 0.14 AND other flags). Do NOT "
    "downgrade just because some metrics are imperfect.\n\n"
    "NOW EVALUATE:\n"
    "{stats_block}\n\n"
    "Think step by step, then give verdict.\n"
    'JSON only: {{"keep_A": true, "reasoning": "1-2 sentences"}}'
)


def _format_stats_for_judge(
    divergence: float,
    n_clipping: int,
    policy_fit: float,
    consistency: float,
    elevated_cvs: int,
    total_cvs: int = 8,
) -> str:
    """Format key metrics for the AI judge prompt."""
    return (
        f"Divergence: {divergence:.3f}, "
        f"Clipping: {n_clipping} joints, "
        f"PolicyFit: {policy_fit:.2f}, "
        f"ElevatedCVs: {elevated_cvs}/{total_cvs}, "
        f"Consistency: {consistency:.3f}"
    )


def check_grade(
    divergence: float,
    n_clipping: int,
    policy_fit: float,
    consistency: float,
    elevated_cvs: int,
) -> AIJudgeResult:
    """Ask AI judge whether an A grade should be kept or downgraded.

    Only call this for A-graded datasets. Returns a recommendation
    with reasoning.

    Requires GOOGLE_API_KEY environment variable.

    Args:
        divergence: Action divergence score (0-1).
        n_clipping: Number of clipping joints.
        policy_fit: Policy fit score (0-1).
        consistency: Overall consistency score (0-1).
        elevated_cvs: Number of CVs > 0.50.

    Returns:
        AIJudgeResult with keep_grade and reasoning.
    """
    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        return AIJudgeResult(
            keep_grade=True,
            reasoning="No API key — keeping current grade",
            model_used="none",
            error="GOOGLE_API_KEY not set",
        )

    stats_block = _format_stats_for_judge(
        divergence, n_clipping, policy_fit, consistency, elevated_cvs,
    )
    prompt = _BINARY_JUDGE_PROMPT.format(stats_block=stats_block)

    try:
        from google import genai

        client = genai.Client(api_key=api_key)
        response = client.models.generate_content(
            model="gemini-2.5-flash",
            contents=[prompt],
        )

        text = response.text.strip()
        if text.startswith("```"):
            lines = text.split("\n")
            text = "\n".join(l for l in lines if not l.strip().startswith("```"))

        import re

        try:
            data = json.loads(text)
        except json.JSONDecodeError:
            match = re.search(r"\{[\s\S]*?\}", text)
            if match:
                data = json.loads(match.group())
            else:
                return AIJudgeResult(
                    keep_grade=True,
                    reasoning="Failed to parse AI response — keeping grade",
                    model_used="gemini-2.5-flash",
                    error=f"JSON parse failed: {text[:200]}",
                )

        return AIJudgeResult(
            keep_grade=bool(data.get("keep_A", True)),
            reasoning=data.get("reasoning", "No reasoning provided"),
            model_used="gemini-2.5-flash",
        )

    except ImportError:
        return AIJudgeResult(
            keep_grade=True,
            reasoning="google-genai not installed",
            model_used="none",
            error="pip install google-genai required",
        )
    except Exception as e:
        logger.warning("AI quality judge failed: %s", e)
        return AIJudgeResult(
            keep_grade=True,
            reasoning=f"AI judge error — keeping grade: {e}",
            model_used="gemini-2.5-flash",
            error=str(e),
        )
