"""GeminiDataAnalyst — AI reasoning layer for robotics data understanding.

Uses Gemini as a reasoning engine for cases where statistics alone aren't
sufficient. Three main capabilities:

1. UNDERSTAND: Given array profiles + metadata, produce a natural language
   description of what the dataset contains and what policy would work best.

2. DIAGNOSE: Given statistical analysis, explain WHY data has issues and
   WHAT to do about them — not just numbers, but actionable explanations.

3. PREDICT: Given the full analysis, produce nuanced training predictions
   with specific risk factors and improvement suggestions.

This extends the existing deep_analysis.py with intelligence from
array_intelligence.py. Falls back gracefully when no API key is available.
"""

from __future__ import annotations

import json
import logging
from dataclasses import asdict, dataclass, field

from orbit.analyzer.array_intelligence import ArrayProfile
from orbit.llm_utils import cache_get, cache_key, cache_set, call_gemini, load_api_key

logger = logging.getLogger(__name__)


@dataclass
class DatasetUnderstanding:
    """Natural language understanding of what a dataset contains."""

    description: str  # "This is a 7-DOF single-arm manipulation dataset..."
    task_description: str  # "Pick-and-place task with a parallel gripper"
    action_space_description: str  # "Absolute joint positions + binary gripper"
    recommended_policy: str  # "ACT with chunk_size=20"
    confidence: float  # 0-1


@dataclass
class DiagnosisReport:
    """Expert diagnosis of data quality issues."""

    summary: str  # "The main issue is inconsistent grasp timing"
    issues: list[dict]  # [{issue, explanation, affected_episodes, fix}]
    overall_assessment: str  # "Trainable with cleanup, expect 70% success"


@dataclass
class TrainingPrediction:
    """Nuanced training outcome prediction."""

    predicted_success_rate: float
    confidence_range: tuple[float, float]
    risk_factors: list[str]
    improvement_suggestions: list[dict]  # [{suggestion, expected_impact}]
    training_tips: list[str]


class GeminiDataAnalyst:
    """AI reasoning layer for robotics data quality analysis.

    Uses Gemini to provide expert-level understanding of robotics datasets
    beyond what statistical analysis can capture.
    """

    def __init__(self, api_key: str | None = None):
        self.api_key = api_key or load_api_key()
        self._available = self.api_key is not None

    @property
    def available(self) -> bool:
        return self._available

    def understand_dataset(
        self,
        array_profiles: dict[str, ArrayProfile],
        metadata: dict | None = None,
    ) -> DatasetUnderstanding | None:
        """Produce a natural language understanding of the dataset.

        Args:
            array_profiles: Dict mapping "action"/"state" to ArrayProfile.
            metadata: Optional metadata dict (repo_id, robot_type, shapes, etc.).

        Returns:
            DatasetUnderstanding, or None if Gemini is unavailable.
        """
        if not self._available:
            return None

        profile_text = self._format_profiles(array_profiles)
        meta_text = json.dumps(metadata, indent=2, default=str) if metadata else "No metadata available"

        prompt = f"""You are an expert robotics ML engineer. Analyze this dataset and explain
what it contains.

ARRAY PROFILES (from statistical analysis):
{profile_text}

METADATA:
{meta_text}

Based on the statistical properties, answer EXACTLY in this JSON format:
{{
    "description": "One paragraph describing the dataset — robot type, DOF, task, format",
    "task_description": "What task the demonstrations show (e.g., pick-and-place, insertion)",
    "action_space_description": "What the action dimensions represent (e.g., 6 joint angles + 1 gripper)",
    "recommended_policy": "Which policy would work best and why (ACT, Diffusion Policy, etc.)",
    "confidence": 0.0-1.0
}}

Be specific. Use the dimension count, value ranges, autocorrelation, and bimodality
to infer the action space. For example:
- 14 dims with bimanual structure + 2 gripper dims = ALOHA-like dual-arm
- 7 dims with 1 gripper dim + delta action type = single-arm EEF control
- 2 dims with no gripper = planar task (like Push-T)"""

        ck = cache_key(prompt)
        cached = cache_get(ck, "gemini_understand")
        if cached:
            return self._parse_understanding(cached)

        response = call_gemini(prompt, self.api_key, max_tokens=1024)
        if not response:
            return None

        cache_set(ck, response, "gemini_understand")
        return self._parse_understanding(response)

    def diagnose_problems(
        self,
        analysis_results: dict,
        array_profiles: dict[str, ArrayProfile] | None = None,
    ) -> DiagnosisReport | None:
        """Explain WHY data has issues and WHAT to do about them.

        Not "divergence is 0.35" but "episodes 12-18 appear to demonstrate
        a different strategy than episodes 1-11."

        Args:
            analysis_results: Full orbit analyze JSON output.
            array_profiles: Optional array profiles for deeper understanding.

        Returns:
            DiagnosisReport, or None if Gemini is unavailable.
        """
        if not self._available:
            return None

        # Format the key analysis results (not the full verbose JSON)
        metrics = self._format_analysis_summary(analysis_results)
        profile_text = self._format_profiles(array_profiles) if array_profiles else ""

        prompt = f"""You are an expert robotics ML engineer diagnosing a training dataset.
Do NOT just restate numbers. Explain what they MEAN for training and what to DO about it.

ANALYSIS RESULTS:
{metrics}

{f"ARRAY PROFILES:{chr(10)}{profile_text}" if profile_text else ""}

For each issue found, explain:
1. WHAT is happening (in terms a roboticist understands)
2. WHY it matters for training (specific failure mode it will cause)
3. WHAT to do about it (specific, actionable steps)

Respond in this JSON format:
{{
    "summary": "One sentence overall assessment",
    "issues": [
        {{
            "issue": "Name of the problem",
            "explanation": "What this means for training — be specific about failure modes",
            "affected_episodes": "Which episodes or dimensions are affected",
            "fix": "Exactly what to do — include orbit commands if applicable"
        }}
    ],
    "overall_assessment": "Will this train successfully? What success rate to expect?"
}}

If the data is clean (grade A/B), say so clearly — don't invent problems."""

        ck = cache_key(prompt)
        cached = cache_get(ck, "gemini_diagnose")
        if cached:
            return self._parse_diagnosis(cached)

        response = call_gemini(prompt, self.api_key, max_tokens=4096)
        if not response:
            return None

        cache_set(ck, response, "gemini_diagnose")
        return self._parse_diagnosis(response)

    def predict_training(
        self,
        analysis_results: dict,
        policy_type: str,
        array_profiles: dict[str, ArrayProfile] | None = None,
    ) -> TrainingPrediction | None:
        """Produce nuanced training outcome prediction.

        Not just "72% success" but "with ACT and 50 episodes, expect ~75% success.
        The main risk is the 3 episodes with gripper timing issues..."

        Args:
            analysis_results: Full orbit analyze JSON output.
            policy_type: Target policy type.
            array_profiles: Optional array profiles.

        Returns:
            TrainingPrediction, or None if Gemini is unavailable.
        """
        if not self._available:
            return None

        metrics = self._format_analysis_summary(analysis_results)
        profile_text = self._format_profiles(array_profiles) if array_profiles else ""

        prompt = f"""You are an expert who has trained 500+ imitation learning policies.
Given this dataset analysis, predict training outcomes for {policy_type}.

ANALYSIS:
{metrics}

{f"ARRAY PROFILES:{chr(10)}{profile_text}" if profile_text else ""}

Respond in JSON:
{{
    "predicted_success_rate": 0.0-1.0,
    "confidence_range": [low, high],
    "risk_factors": ["specific risk 1", "specific risk 2"],
    "improvement_suggestions": [
        {{"suggestion": "specific action", "expected_impact": "+X% success rate"}}
    ],
    "training_tips": [
        "Loss should drop below X by step Y",
        "Watch for specific behavior at eval"
    ]
}}

Be precise. Reference specific episodes, dimensions, or metrics.
Base predictions on published results (RoboMimic, ALOHA, Push-T baselines)."""

        ck = cache_key(prompt)
        cached = cache_get(ck, "gemini_predict")
        if cached:
            return self._parse_prediction(cached)

        response = call_gemini(prompt, self.api_key, max_tokens=1536)
        if not response:
            return None

        cache_set(ck, response, "gemini_predict")
        return self._parse_prediction(response)

    # ------------------------------------------------------------------
    # Formatting helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _format_profiles(profiles: dict[str, ArrayProfile] | None) -> str:
        if not profiles:
            return "No profiles available"

        lines = []
        for name, prof in profiles.items():
            lines.append(f"\n{name.upper()} ARRAY ({prof.num_dimensions} dimensions):")
            lines.append(f"  Action type: {prof.action_type} (confidence: {prof.action_type_confidence:.2f})")
            lines.append(f"  Normalized: {prof.is_normalized}, Scale: {prof.scale:.4f}")
            lines.append(f"  Smoothness: {prof.smoothness:.3f}")
            lines.append(f"  Bimanual: {prof.is_bimanual}")

            if prof.dead_dimensions:
                lines.append(f"  Dead dims: {prof.dead_dimensions}")
            if prof.gripper_dimensions:
                lines.append(f"  Gripper dims: {prof.gripper_dimensions}")
            if prof.clipping_dimensions:
                lines.append(f"  Clipping dims: {prof.clipping_dimensions}")
            if prof.noisy_dimensions:
                lines.append(f"  Noisy dims: {prof.noisy_dimensions}")

            lines.append("  Per-dimension roles:")
            for dp in prof.dimension_profiles[:20]:  # cap for context limits
                lines.append(
                    f"    dim[{dp.index}]: {dp.role.role} (conf={dp.role.confidence:.2f}) "
                    f"range=[{dp.min_val:.4f}, {dp.max_val:.4f}] "
                    f"autocorr={dp.lag1_autocorrelation:.3f}"
                )

        return "\n".join(lines)

    @staticmethod
    def _format_analysis_summary(analysis: dict) -> str:
        lines = []
        lines.append(f"Dataset: {analysis.get('repo_id', 'unknown')}")
        lines.append(f"Episodes: {analysis.get('num_episodes', '?')}")
        lines.append(f"Frames: {analysis.get('num_frames', '?')}")

        grade = analysis.get("readiness", {})
        if grade:
            lines.append(f"Grade: {grade.get('grade', '?')} ({grade.get('score', '?')}/100)")
            for p in grade.get("problems", []):
                lines.append(f"  Problem: {p}")
            for s in grade.get("strengths", []):
                lines.append(f"  Strength: {s}")

        cov = analysis.get("coverage", {})
        if cov:
            lines.append(f"Coverage: consistency={cov.get('episode_consistency', '?')}")

        ad = analysis.get("action_divergence", {})
        if ad:
            lines.append(f"Action divergence: {ad.get('score', '?')}")

        sd = analysis.get("signal_diagnostics", {})
        if sd:
            dead = sd.get("dead_dimensions", [])
            if dead:
                lines.append(f"Dead joints: {dead}")

        quality = analysis.get("quality_score")
        if quality is not None:
            lines.append(f"Quality score: {quality:.3f}")

        return "\n".join(lines)

    # ------------------------------------------------------------------
    # Response parsing helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _parse_understanding(response: str) -> DatasetUnderstanding | None:
        try:
            data = _extract_json(response)
            return DatasetUnderstanding(
                description=data.get("description", ""),
                task_description=data.get("task_description", ""),
                action_space_description=data.get("action_space_description", ""),
                recommended_policy=data.get("recommended_policy", ""),
                confidence=float(data.get("confidence", 0.5)),
            )
        except Exception as e:
            logger.warning("Failed to parse understanding response: %s", e)
            return DatasetUnderstanding(
                description=response[:500],
                task_description="",
                action_space_description="",
                recommended_policy="",
                confidence=0.3,
            )

    @staticmethod
    def _parse_diagnosis(response: str) -> DiagnosisReport | None:
        try:
            data = _extract_json(response)
            return DiagnosisReport(
                summary=data.get("summary", ""),
                issues=data.get("issues", []),
                overall_assessment=data.get("overall_assessment", ""),
            )
        except Exception as e:
            logger.warning("Failed to parse diagnosis response: %s", e)
            return DiagnosisReport(
                summary=response[:300],
                issues=[],
                overall_assessment="",
            )

    @staticmethod
    def _parse_prediction(response: str) -> TrainingPrediction | None:
        try:
            data = _extract_json(response)
            cr = data.get("confidence_range", [0.3, 0.8])
            return TrainingPrediction(
                predicted_success_rate=float(data.get("predicted_success_rate", 0.5)),
                confidence_range=(float(cr[0]), float(cr[1])),
                risk_factors=data.get("risk_factors", []),
                improvement_suggestions=data.get("improvement_suggestions", []),
                training_tips=data.get("training_tips", []),
            )
        except Exception as e:
            logger.warning("Failed to parse prediction response: %s", e)
            return None


def _extract_json(text: str) -> dict:
    """Extract a JSON object from a response that may contain markdown fences."""
    import re

    # Try to find JSON in code fences
    match = re.search(r"```(?:json)?\s*\n?(.*?)\n?```", text, re.DOTALL)
    if match:
        return json.loads(match.group(1))

    # Try to find a JSON object directly
    start = text.find("{")
    end = text.rfind("}") + 1
    if start >= 0 and end > start:
        return json.loads(text[start:end])

    raise ValueError(f"No JSON found in response: {text[:200]}")
