"""VLM-as-predictor — replaces hand-tuned quality formula with a VLM that reasons
about dataset quality using ground truth benchmarks as context.

Supports Gemini Flash (cheapest, default) and Anthropic Claude as fallback.
Falls back to heuristic scoring when no API key is available.
"""

from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.signal_diagnostics import SignalDiagnosticsReport

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Ground truth cache (loaded once)
# ---------------------------------------------------------------------------

_ground_truth_cache: list[dict] | None = None

GROUND_TRUTH_PATH = Path(__file__).parent / "data" / "ground_truth_condensed.json"


def _load_ground_truth() -> list[dict]:
    """Load ground truth data, caching after first call."""
    global _ground_truth_cache
    if _ground_truth_cache is not None:
        return _ground_truth_cache

    try:
        with open(GROUND_TRUTH_PATH) as f:
            _ground_truth_cache = json.load(f)
        logger.info("Loaded %d ground truth entries", len(_ground_truth_cache))
    except (FileNotFoundError, json.JSONDecodeError) as e:
        logger.warning("Failed to load ground truth data: %s", e)
        _ground_truth_cache = []

    return _ground_truth_cache


def _select_relevant_entries(
    ground_truth: list[dict],
    task_type: str | None = None,
    policy: str | None = None,
    n: int = 10,
) -> list[dict]:
    """Select the N most relevant ground truth entries for the prompt.

    Prioritizes entries matching task_type, then policy, then fills with
    diverse entries.
    """
    if not ground_truth:
        return []

    scored: list[tuple[float, dict]] = []
    task_type_lower = (task_type or "").lower()
    policy_lower = (policy or "").lower()

    for entry in ground_truth:
        score = 0.0
        entry_task = entry.get("task_type", "").lower()
        entry_policy = entry.get("policy", "").lower()
        entry_summary = entry.get("features_summary", "").lower()

        # Task type matching
        if task_type_lower and task_type_lower in entry_task:
            score += 3.0
        elif task_type_lower and any(
            w in entry_task or w in entry_summary for w in task_type_lower.split("_")
        ):
            score += 1.5

        # Policy matching
        if policy_lower and policy_lower in entry_policy:
            score += 2.0
        elif policy_lower and any(w in entry_policy for w in policy_lower.split("_") if len(w) > 2):
            score += 1.0

        # Slight randomness for diversity
        score += hash(entry.get("id", "")) % 100 / 1000.0

        scored.append((score, entry))

    scored.sort(key=lambda x: x[0], reverse=True)
    return [entry for _, entry in scored[:n]]


# ---------------------------------------------------------------------------
# Grade mapping
# ---------------------------------------------------------------------------


def _grade_from_score(score: float) -> str:
    """Map a predicted success rate to a letter grade."""
    if score >= 0.80:
        return "A"
    elif score >= 0.65:
        return "B"
    elif score >= 0.50:
        return "C"
    elif score >= 0.35:
        return "D"
    else:
        return "F"


# ---------------------------------------------------------------------------
# Prompt construction
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = (
    "You are an expert robotics researcher evaluating robot training datasets. "
    "You have access to 82 validated benchmark results from published papers. "
    "Given a dataset's statistics, predict its likely success rate and explain why. "
    "\n\n"
    "IMPORTANT CONTEXT ON DIRECTIONAL BIAS: In goal-directed robot manipulation "
    "(pick-and-place, insertion, transfer, pushing, lifting), it is NORMAL and EXPECTED "
    "for most joints to show directional bias — the robot naturally moves toward the object, "
    "grasps it, and moves to the target. When the signal diagnostics show directional bias "
    "warnings on many joints simultaneously, this almost always indicates goal-directed motion, "
    "NOT a data quality problem. Treat widespread directional bias as informational, not as a "
    "quality defect. Only flag directional bias as a real concern when (a) a single isolated "
    "joint is biased while others are not (suggesting a hardware or calibration issue), or "
    "(b) the dataset is explicitly intended for random exploration or coverage tasks."
    "\n\n"
    "Respond ONLY with a JSON object matching this exact schema: "
    '{"predicted_success_rate": <float 0-1>, "confidence": <float 0-1>, '
    '"grade": "<A/B/C/D/F>", "reasoning": "<2-3 sentences>", '
    '"similar_datasets": [{"id": "...", "success_rate": ..., "why_similar": "..."}], '
    '"key_strengths": ["..."], "key_weaknesses": ["..."], '
    '"improvement_estimate": "<string>"}'
)


def _build_prompt(
    coverage_report: FullCoverageReport,
    signal_diagnostics: SignalDiagnosticsReport | None,
    quality_metrics: dict,
    dataset_info: dict,
    reference_entries: list[dict],
) -> str:
    """Build the user prompt with all dataset statistics and reference data."""
    lines = [
        "# Dataset Under Evaluation",
        "",
        f"Task: {dataset_info.get('task_description', 'unknown')}",
        f"Robot: {dataset_info.get('robot_type', 'unknown')}",
        f"Policy: {dataset_info.get('policy', 'unknown')}",
        f"Episodes: {dataset_info.get('num_episodes', 'unknown')}",
        "",
        "## Coverage Scores (0=none, 1=perfect)",
        (
            f"- position_diversity: "
            f"{coverage_report.position_diversity:.3f} — "
            + (
                "good"
                if coverage_report.position_diversity >= 0.65
                else f"the robot only visits "
                f"{int(coverage_report.position_diversity * 100)}"
                f"% of the workspace"
            )
        ),
        f"- action_diversity: {coverage_report.action_diversity:.3f}",
        f"- episode_consistency: {coverage_report.episode_consistency:.3f}",
        f"- temporal_coverage: {coverage_report.temporal_coverage:.3f}",
    ]

    # Quality metrics
    if quality_metrics:
        lines.extend(["", "## Quality Metrics"])
        for k, v in quality_metrics.items():
            if isinstance(v, float):
                lines.append(f"- {k}: {v:.3f}")
            else:
                lines.append(f"- {k}: {v}")

    # Signal diagnostics
    if signal_diagnostics is not None:
        lines.extend(["", "## Signal Diagnostics"])

        if signal_diagnostics.blockers:
            lines.append("BLOCKERS:")
            for b in signal_diagnostics.blockers:
                lines.append(f"  - {b}")

        if signal_diagnostics.warnings:
            lines.append("WARNINGS:")
            for w in signal_diagnostics.warnings:
                lines.append(f"  - {w}")

        if signal_diagnostics.passing:
            lines.append("PASSING:")
            for p in signal_diagnostics.passing:
                lines.append(f"  - {p}")

        if signal_diagnostics.prescription:
            lines.append("PRESCRIPTION:")
            for rx in signal_diagnostics.prescription:
                lines.append(f"  - {rx}")

        lines.append(f"Dead dimensions: {signal_diagnostics.dead_dimensions}")
        lines.append(f"Episodes to remove: {signal_diagnostics.episodes_to_remove}")

    # Reference data
    if reference_entries:
        lines.extend(
            [
                "",
                "## Reference Benchmarks (from 82 published results)",
                "Use these to calibrate your prediction:",
            ]
        )
        for entry in reference_entries:
            lines.append(
                f"- {entry['id']}: {entry['dataset']} | "
                f"policy={entry['policy']} | "
                f"success={entry['success_rate']} | "
                f"eps={entry['num_episodes']} | "
                f"{entry.get('features_summary', '')}"
            )

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# VLM API calls
# ---------------------------------------------------------------------------


def _call_gemini(
    system: str,
    user: str,
    api_key: str,
    frames: list[str] | None = None,
) -> tuple[dict, int, float]:
    """Call Google Gemini Flash and return (parsed_json, tokens, cost)."""
    import google.generativeai as genai

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

    # Build content parts
    parts: list[Any] = []
    if frames:
        for i, frame_b64 in enumerate(frames[:5]):  # max 5 frames
            parts.append(
                {
                    "inline_data": {
                        "mime_type": "image/jpeg",
                        "data": frame_b64,
                    }
                }
            )
        parts.append(f"Above are {len(parts)} sample frames from the dataset.\n\n{user}")
    else:
        parts.append(user)

    response = model.generate_content(parts)
    text = response.text.strip()
    parsed = _parse_vlm_json(text)

    tokens_meta = getattr(response, "usage_metadata", None)
    token_count = getattr(tokens_meta, "total_token_count", 0) if tokens_meta else 0
    # Gemini Flash: ~$0.003/call
    cost = max(0.001, token_count * 0.000001)

    return parsed, token_count, cost


def _call_anthropic(
    system: str,
    user: str,
    api_key: str,
    frames: list[str] | None = None,
) -> tuple[dict, int, float]:
    """Call Anthropic Claude and return (parsed_json, tokens, cost)."""
    import anthropic

    client = anthropic.Anthropic(api_key=api_key)

    # Build content
    content: list[dict] = []
    if frames:
        for frame_b64 in frames[:5]:
            content.append(
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/jpeg",
                        "data": frame_b64,
                    },
                }
            )
    content.append({"type": "text", "text": user})

    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=1024,
        system=system,
        messages=[{"role": "user", "content": content}],
    )

    text = response.content[0].text.strip()
    parsed = _parse_vlm_json(text)

    input_tokens = response.usage.input_tokens
    output_tokens = response.usage.output_tokens
    total_tokens = input_tokens + output_tokens
    # Claude Sonnet: $3/M input, $15/M output
    cost = (input_tokens * 3 + output_tokens * 15) / 1_000_000

    return parsed, total_tokens, cost


def _parse_vlm_json(text: str) -> dict:
    """Robustly parse VLM response — handles markdown code blocks, partial JSON."""
    # Strip markdown code blocks
    if "```" in text:
        lines = text.split("\n")
        in_block = False
        json_lines = []
        for line in lines:
            if line.strip().startswith("```"):
                if in_block:
                    break
                in_block = True
                continue
            if in_block:
                json_lines.append(line)
        if json_lines:
            text = "\n".join(json_lines)

    # Try parsing directly
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # Try finding JSON object in text
    start = text.find("{")
    end = text.rfind("}")
    if start != -1 and end != -1 and end > start:
        try:
            return json.loads(text[start : end + 1])
        except json.JSONDecodeError:
            pass

    # Log raw response for debugging
    logger.warning("Failed to parse VLM response as JSON. Raw response:\n%s", text[:500])
    raise json.JSONDecodeError("Could not parse VLM response", text, 0)


# ---------------------------------------------------------------------------
# Heuristic fallback
# ---------------------------------------------------------------------------


def _heuristic_prediction(
    coverage_report: FullCoverageReport,
    quality_metrics: dict,
    dataset_info: dict,
    signal_diagnostics: SignalDiagnosticsReport | None = None,
) -> VLMPrediction:
    """Fall back to heuristic quality scoring when no VLM is available."""
    # Use the enhanced quality formula
    coverage_score = (
        coverage_report.position_diversity * 0.25
        + coverage_report.action_diversity * 0.20
        + coverage_report.episode_consistency * 0.15
        + coverage_report.temporal_coverage * 0.15
    )
    episode_count = dataset_info.get("num_episodes", 0)
    episode_score = min(1.0, episode_count / 50) * 0.25

    base_score = coverage_score + episode_score

    # Incorporate quality metrics if available
    smoothness = quality_metrics.get("action_smoothness", 0.5)
    mi = quality_metrics.get("mutual_information", 0.5)
    action_quality = smoothness * 0.5 + mi * 0.5

    # Signal diagnostics penalty
    signal_penalty = 0.0
    if signal_diagnostics is not None:
        signal_penalty = len(signal_diagnostics.blockers) * 0.05
        signal_penalty += len(signal_diagnostics.warnings) * 0.02

    predicted = float(np.clip(base_score * 0.6 + action_quality * 0.25 - signal_penalty, 0.0, 1.0))

    # Build strengths/weaknesses
    strengths = []
    weaknesses = []
    if coverage_report.position_diversity >= 0.65:
        strengths.append("Good workspace coverage")
    else:
        weaknesses.append(f"Low position diversity ({coverage_report.position_diversity:.2f})")

    if coverage_report.action_diversity >= 0.65:
        strengths.append("Good action diversity")
    else:
        weaknesses.append(f"Low action diversity ({coverage_report.action_diversity:.2f})")

    if episode_count >= 50:
        strengths.append(f"Sufficient episodes ({episode_count})")
    else:
        weaknesses.append(f"Low episode count ({episode_count})")

    if signal_diagnostics and signal_diagnostics.blockers:
        for b in signal_diagnostics.blockers:
            weaknesses.append(b)

    grade = _grade_from_score(predicted)

    return VLMPrediction(
        predicted_success_rate=predicted,
        confidence=0.4,  # Low confidence for heuristic
        grade=grade,
        reasoning=(
            f"Heuristic estimate based on coverage ({base_score:.2f}), "
            f"action quality ({action_quality:.2f}), and "
            f"{len(signal_diagnostics.blockers) if signal_diagnostics else 0} blockers. "
            "Set GOOGLE_API_KEY or ANTHROPIC_API_KEY for VLM-powered prediction."
        ),
        similar_datasets=[],
        key_strengths=strengths[:3],
        key_weaknesses=weaknesses[:3],
        improvement_estimate="Unable to estimate without VLM analysis",
        model_used="heuristic",
        tokens_used=0,
        cost_estimate=0.0,
    )


# ---------------------------------------------------------------------------
# Dataclass
# ---------------------------------------------------------------------------


@dataclass
class VLMPrediction:
    """Result of VLM-based quality prediction."""

    predicted_success_rate: float  # 0.0 to 1.0
    confidence: float  # 0.0 to 1.0
    grade: str  # A/B/C/D/F
    reasoning: str  # VLM's explanation
    similar_datasets: list[dict] = field(default_factory=list)
    key_strengths: list[str] = field(default_factory=list)
    key_weaknesses: list[str] = field(default_factory=list)
    improvement_estimate: str = ""
    model_used: str = ""
    tokens_used: int = 0
    cost_estimate: float = 0.0


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------


def predict_success_rate(
    coverage_report: FullCoverageReport,
    signal_diagnostics: SignalDiagnosticsReport | None = None,
    quality_metrics: dict | None = None,
    dataset_info: dict | None = None,
    sample_frames_base64: list[str] | None = None,
    api_key: str | None = None,
    anthropic_key: str | None = None,
) -> VLMPrediction:
    """Predict dataset success rate using a VLM with ground truth context.

    Tries Gemini first (cheaper), falls back to Anthropic Claude, then to
    heuristic if no API key is available. Never crashes.

    Args:
        coverage_report: Full coverage analysis results.
        signal_diagnostics: Optional signal diagnostics report.
        quality_metrics: Dict with action_smoothness, mutual_information, etc.
        dataset_info: Dict with num_episodes, task_description, robot_type, policy.
        sample_frames_base64: Optional base64-encoded frames for vision analysis.
        api_key: Gemini API key (falls back to GOOGLE_API_KEY env var).
        anthropic_key: Claude API key (falls back to ANTHROPIC_API_KEY env var).

    Returns:
        VLMPrediction with predicted success rate and analysis.
    """
    quality_metrics = quality_metrics or {}
    dataset_info = dataset_info or {}

    # Resolve API keys
    gemini_key = api_key or os.environ.get("GOOGLE_API_KEY")
    claude_key = anthropic_key or os.environ.get("ANTHROPIC_API_KEY")

    # If no keys, go straight to heuristic
    if not gemini_key and not claude_key:
        logger.info("No API keys available, using heuristic fallback")
        return _heuristic_prediction(
            coverage_report, quality_metrics, dataset_info, signal_diagnostics
        )

    # Load ground truth and select relevant entries
    ground_truth = _load_ground_truth()
    reference_entries = _select_relevant_entries(
        ground_truth,
        task_type=dataset_info.get("task_description"),
        policy=dataset_info.get("policy"),
        n=10,
    )

    # Build prompt
    user_prompt = _build_prompt(
        coverage_report=coverage_report,
        signal_diagnostics=signal_diagnostics,
        quality_metrics=quality_metrics,
        dataset_info=dataset_info,
        reference_entries=reference_entries,
    )

    # Try Gemini first
    if gemini_key:
        try:
            logger.info("Calling Gemini Flash for prediction")
            parsed, tokens, cost = _call_gemini(
                SYSTEM_PROMPT, user_prompt, gemini_key, sample_frames_base64
            )
            return _build_prediction(parsed, "gemini-2.5-flash", tokens, cost)
        except ImportError:
            logger.warning("google-generativeai not installed, trying Claude")
        except Exception as e:
            logger.warning("Gemini call failed: %s, trying Claude", e)

    # Try Claude
    if claude_key:
        try:
            logger.info("Calling Claude for prediction")
            parsed, tokens, cost = _call_anthropic(
                SYSTEM_PROMPT, user_prompt, claude_key, sample_frames_base64
            )
            return _build_prediction(parsed, "claude-sonnet", tokens, cost)
        except ImportError:
            logger.warning("anthropic SDK not installed, falling back to heuristic")
        except Exception as e:
            logger.warning("Claude call failed: %s, falling back to heuristic", e)

    # Final fallback
    logger.info("All VLM calls failed, using heuristic fallback")
    return _heuristic_prediction(coverage_report, quality_metrics, dataset_info, signal_diagnostics)


def _build_prediction(
    parsed: dict,
    model: str,
    tokens: int,
    cost: float,
) -> VLMPrediction:
    """Build VLMPrediction from parsed VLM JSON response, filling defaults."""
    predicted = float(parsed.get("predicted_success_rate", 0.5))
    predicted = max(0.0, min(1.0, predicted))

    confidence = float(parsed.get("confidence", 0.5))
    confidence = max(0.0, min(1.0, confidence))

    grade = parsed.get("grade", _grade_from_score(predicted))

    similar = parsed.get("similar_datasets", [])
    if not isinstance(similar, list):
        similar = []

    strengths = parsed.get("key_strengths", [])
    if not isinstance(strengths, list):
        strengths = []

    weaknesses = parsed.get("key_weaknesses", [])
    if not isinstance(weaknesses, list):
        weaknesses = []

    return VLMPrediction(
        predicted_success_rate=predicted,
        confidence=confidence,
        grade=grade,
        reasoning=parsed.get("reasoning", "No reasoning provided"),
        similar_datasets=similar,
        key_strengths=strengths,
        key_weaknesses=weaknesses,
        improvement_estimate=parsed.get("improvement_estimate", ""),
        model_used=model,
        tokens_used=tokens,
        cost_estimate=cost,
    )
