"""LLM-powered recommendations — send statistics to an LLM for insightful analysis.

Supports Google Gemini (cheapest, default) and Anthropic Claude.
Falls back to template recommendations when no API key is available.
"""

from __future__ import annotations

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

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.dataset_loader import DatasetInfo
from orbit.analyzer.task_inference import TaskProfile

logger = logging.getLogger(__name__)


@dataclass
class Recommendation:
    """A single actionable recommendation."""

    priority: str  # "HIGH", "MEDIUM", "LOW"
    action: str
    reasoning: str
    expected_impact: float = 0.0
    phase_affected: str = ""


@dataclass
class LLMRecommendationResult:
    """Complete result from LLM recommendation generation."""

    overall_assessment: str = ""
    key_insight: str = ""
    recommendations: list[Recommendation] = field(default_factory=list)
    model_used: str = ""
    tokens_used: int = 0
    cost_usd: float = 0.0


SYSTEM_PROMPT = """You are an expert robotics researcher specializing in imitation learning \
data quality. You interpret statistical analysis of robot datasets and produce \
specific, actionable recommendations.

Rules:
- Be specific: "Collect 12 episodes with side grasps" not "collect more data"
- Reference actual numbers from the statistics
- Explain WHY each recommendation matters for THIS task
- Connect statistics to likely real-world failure modes
- If episode durations suggest truncated demos, flag this prominently
- If coverage is good everywhere, say so — don't invent problems
- Maximum 6 recommendations ranked by impact

Respond in JSON:
{
  "overall_assessment": "...",
  "key_insight": "...",
  "recommendations": [
    {
      "priority": "HIGH|MEDIUM|LOW",
      "action": "specific instruction",
      "reasoning": "why this matters for this task",
      "expected_impact": 0.09,
      "phase_affected": "grasp"
    }
  ]
}"""


def build_user_prompt(
    info: DatasetInfo,
    coverage: FullCoverageReport,
    task_profile: TaskProfile,
    task_description: str = "",
    embedding_diversity_score: float | None = None,
) -> str:
    """Build the user prompt with all statistics."""
    dur_lo, dur_hi = task_profile.expected_episode_duration_range

    mean_duration = "unknown"
    if info.num_episodes > 0 and info.fps > 0:
        mean_duration = f"{info.num_frames / info.num_episodes / info.fps:.1f}"

    lines = [
        f'TASK: {task_profile.task_type} — "{task_description}"',
        f"ROBOT: {info.robot_type or 'unknown'}",
        "",
        "STATISTICS:",
        f"- Episodes: {info.num_episodes} (recommended: {task_profile.min_recommended_episodes})",
        f"- Mean episode duration: {mean_duration}s (expected: {dur_lo}-{dur_hi}s)",
        f"- Frames: {info.num_frames} | FPS: {info.fps}",
        "",
        "COVERAGE (0=none, 1=perfect):",
        f"- Position diversity: {coverage.position_diversity:.2f}",
        f"- Action diversity: {coverage.action_diversity:.2f}",
        f"- Episode consistency: {coverage.episode_consistency:.2f}",
        f"- Temporal coverage: {coverage.temporal_coverage:.2f}",
    ]

    if embedding_diversity_score is not None:
        lines.append(f"- Visual diversity: {embedding_diversity_score:.2f}")

    lines.extend(
        [
            "",
            "KNOWN FAILURE MODES FOR THIS TASK:",
        ]
    )
    for fm in task_profile.typical_failure_modes:
        lines.append(f"- {fm}")

    return "\n".join(lines)


def _call_gemini(system: str, user: str, api_key: str) -> tuple[dict, str, int, float]:
    """Call Google Gemini Flash."""
    import google.generativeai as genai

    genai.configure(api_key=api_key)
    model = genai.GenerativeModel("gemini-2.5-flash", system_instruction=system)
    response = model.generate_content(user)

    text = response.text.strip()
    if text.startswith("```"):
        text = text.split("\n", 1)[1] if "\n" in text else text[3:]
        if text.endswith("```"):
            text = text[:-3]
        if text.startswith("json"):
            text = text[4:].strip()

    parsed = json.loads(text)
    tokens = getattr(response, "usage_metadata", None)
    token_count = getattr(tokens, "total_token_count", 0) if tokens else 0

    return parsed, "gemini-2.5-flash", token_count, token_count * 0.000001


def _parse_llm_response(data: dict) -> tuple[str, str, list[Recommendation]]:
    """Parse JSON response into structured recommendations."""
    assessment = data.get("overall_assessment", "")
    key_insight = data.get("key_insight", "")

    recs = []
    for r in data.get("recommendations", []):
        recs.append(
            Recommendation(
                priority=r.get("priority", "MEDIUM"),
                action=r.get("action", ""),
                reasoning=r.get("reasoning", ""),
                expected_impact=float(r.get("expected_impact", 0.0)),
                phase_affected=r.get("phase_affected", ""),
            )
        )

    return assessment, key_insight, recs


def generate_llm_recommendations(
    info: DatasetInfo,
    coverage: FullCoverageReport,
    task_profile: TaskProfile,
    task_description: str = "",
    embedding_diversity_score: float | None = None,
    api_key: str | None = None,
    provider: str | None = None,
) -> LLMRecommendationResult:
    """Generate LLM-powered recommendations from dataset statistics."""
    if api_key is None and provider is None:
        if os.environ.get("GOOGLE_API_KEY"):
            provider = "google"
            api_key = os.environ["GOOGLE_API_KEY"]
        else:
            return LLMRecommendationResult(
                overall_assessment=(
                    "No API key found. Set GOOGLE_API_KEY for smart recommendations."
                ),
                model_used="none",
            )
    elif api_key is None:
        env_key = "GOOGLE_API_KEY"
        api_key = os.environ.get(env_key, "")
        if not api_key:
            return LLMRecommendationResult(
                overall_assessment=f"No {env_key} found in environment.",
                model_used="none",
            )

    if provider is None:
        provider = "google"

    user_prompt = build_user_prompt(
        info=info,
        coverage=coverage,
        task_profile=task_profile,
        task_description=task_description,
        embedding_diversity_score=embedding_diversity_score,
    )

    try:
        data, model, tokens, cost = _call_gemini(SYSTEM_PROMPT, user_prompt, api_key)
        assessment, key_insight, recs = _parse_llm_response(data)

        return LLMRecommendationResult(
            overall_assessment=assessment,
            key_insight=key_insight,
            recommendations=recs,
            model_used=model,
            tokens_used=tokens,
            cost_usd=cost,
        )

    except ImportError as exc:
        logger.warning("LLM provider SDK not installed: %s", exc)
        return LLMRecommendationResult(
            overall_assessment="LLM SDK not installed. Install google-generativeai.",
            model_used="none",
        )
    except json.JSONDecodeError as exc:
        logger.warning("Failed to parse LLM response as JSON: %s", exc)
        return LLMRecommendationResult(
            overall_assessment="LLM returned non-JSON response.",
            model_used=provider or "unknown",
        )
    except Exception as exc:
        logger.warning("LLM recommendation failed: %s", exc)
        return LLMRecommendationResult(
            overall_assessment=f"LLM analysis failed: {exc}",
            model_used=provider or "unknown",
        )
