"""Tests for the VLM predictor module."""

import json
from unittest.mock import MagicMock, patch

import numpy as np
import pytest

from orbit.analyzer.coverage import FullCoverageReport
from orbit.analyzer.signal_diagnostics import (
    DirectionalBalance,
    EpisodeDiagnostic,
    GripperAnalysis,
    JointDiagnostic,
    SignalDiagnosticsReport,
)
from orbit.analyzer.vlm_predictor import (
    GROUND_TRUTH_PATH,
    VLMPrediction,
    _build_prediction,
    _grade_from_score,
    _heuristic_prediction,
    _load_ground_truth,
    _parse_vlm_json,
    _select_relevant_entries,
    _build_prompt,
    predict_success_rate,
)


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def coverage_report():
    return FullCoverageReport(
        position_diversity=0.65,
        action_diversity=0.55,
        episode_consistency=0.80,
        temporal_coverage=0.45,
    )


@pytest.fixture
def signal_diagnostics():
    return SignalDiagnosticsReport(
        joint_diagnostics=[
            JointDiagnostic(
                joint_index=0, joint_name="joint_0",
                range_used=(-1.0, 1.0), range_limit=None,
                range_utilization=0.8, zero_velocity_ratio=0.1,
                smoothness=0.7, directional_bias=0.55,
                is_dead=False, is_clipping=False,
            ),
        ],
        gripper_analysis=GripperAnalysis(
            gripper_dim=5, open_ratio=0.6, close_ratio=0.4,
            transitions_per_episode=2.0, has_sufficient_grasps=True,
        ),
        episode_diagnostics=[
            EpisodeDiagnostic(
                episode_index=0, quality_score=0.8,
                length=100, is_truncated=False,
                jerk_spike_count=0,
            ),
        ],
        directional_balance=DirectionalBalance(
            per_joint_bias=[0.55], overall_bias=0.55,
            dominant_direction=None,
        ),
        global_zero_velocity_ratio=0.15,
        dead_dimensions=[],
        episodes_to_remove=[],
        episodes_to_review=[],
        blockers=[],
        warnings=["Episode count (15) below minimum (30)"],
        passing=["All joints active"],
        prescription=["Record 15 more episodes"],
    )


@pytest.fixture
def quality_metrics():
    return {
        "action_smoothness": 0.75,
        "mutual_information": 0.60,
        "action_autocorrelation": 0.70,
    }


@pytest.fixture
def dataset_info():
    return {
        "num_episodes": 50,
        "task_description": "pick up cups from table",
        "robot_type": "so100",
        "policy": "act",
    }


@pytest.fixture
def valid_vlm_response():
    return {
        "predicted_success_rate": 0.62,
        "confidence": 0.75,
        "grade": "B",
        "reasoning": "Good coverage but low temporal diversity. Similar to ALOHA cube transfer.",
        "similar_datasets": [
            {"id": "aloha_transfer_cube_human", "success_rate": 0.82, "why_similar": "Similar manipulation task"},
        ],
        "key_strengths": ["Good workspace coverage", "Sufficient episodes"],
        "key_weaknesses": ["Low temporal coverage", "Moderate action diversity"],
        "improvement_estimate": "Adding 20 more diverse episodes could improve from ~62% to ~75%",
    }


# ---------------------------------------------------------------------------
# Test 1: Heuristic fallback (no API keys)
# ---------------------------------------------------------------------------

def test_heuristic_fallback(coverage_report, quality_metrics, dataset_info):
    """No API keys → falls back to heuristic, returns valid VLMPrediction."""
    with patch.dict("os.environ", {}, clear=True):
        # Remove any API keys
        env = {k: v for k, v in __import__("os").environ.items()
               if k not in ("GOOGLE_API_KEY", "ANTHROPIC_API_KEY")}
        with patch.dict("os.environ", env, clear=True):
            result = predict_success_rate(
                coverage_report=coverage_report,
                quality_metrics=quality_metrics,
                dataset_info=dataset_info,
            )

    assert isinstance(result, VLMPrediction)
    assert result.model_used == "heuristic"
    assert 0.0 <= result.predicted_success_rate <= 1.0
    assert 0.0 <= result.confidence <= 1.0
    assert result.grade in ("A", "B", "C", "D", "F")
    assert result.reasoning != ""
    assert isinstance(result.key_strengths, list)
    assert isinstance(result.key_weaknesses, list)
    assert result.tokens_used == 0
    assert result.cost_estimate == 0.0


# ---------------------------------------------------------------------------
# Test 2: Gemini success
# ---------------------------------------------------------------------------

def test_gemini_success(coverage_report, dataset_info, valid_vlm_response):
    """Mock Gemini to return valid JSON → parsed correctly."""
    mock_response = MagicMock()
    mock_response.text = json.dumps(valid_vlm_response)
    mock_response.usage_metadata = MagicMock()
    mock_response.usage_metadata.total_token_count = 500

    with patch("orbit.analyzer.vlm_predictor._call_gemini") as mock_gemini:
        mock_gemini.return_value = (valid_vlm_response, 500, 0.003)

        result = predict_success_rate(
            coverage_report=coverage_report,
            dataset_info=dataset_info,
            api_key="fake-gemini-key",
        )

    assert isinstance(result, VLMPrediction)
    assert result.model_used == "gemini-2.5-flash"
    assert result.predicted_success_rate == 0.62
    assert result.confidence == 0.75
    assert result.grade == "B"
    assert "coverage" in result.reasoning.lower() or "similar" in result.reasoning.lower() or len(result.reasoning) > 0
    assert len(result.similar_datasets) > 0
    assert result.tokens_used == 500
    assert result.cost_estimate > 0


# ---------------------------------------------------------------------------
# Test 3: Gemini failure → Claude fallback
# ---------------------------------------------------------------------------

def test_gemini_failure_claude_fallback(coverage_report, dataset_info, valid_vlm_response):
    """Gemini raises exception, Claude returns valid JSON → uses Claude."""
    with patch("orbit.analyzer.vlm_predictor._call_gemini") as mock_gemini, \
         patch("orbit.analyzer.vlm_predictor._call_anthropic") as mock_claude:

        mock_gemini.side_effect = RuntimeError("Gemini API error")
        mock_claude.return_value = (valid_vlm_response, 800, 0.01)

        result = predict_success_rate(
            coverage_report=coverage_report,
            dataset_info=dataset_info,
            api_key="fake-gemini-key",
            anthropic_key="fake-claude-key",
        )

    assert result.model_used == "claude-sonnet"
    assert result.predicted_success_rate == 0.62


# ---------------------------------------------------------------------------
# Test 4: Both fail → heuristic
# ---------------------------------------------------------------------------

def test_both_fail_heuristic(coverage_report, quality_metrics, dataset_info):
    """Both APIs fail → falls back to heuristic."""
    with patch("orbit.analyzer.vlm_predictor._call_gemini") as mock_gemini, \
         patch("orbit.analyzer.vlm_predictor._call_anthropic") as mock_claude:

        mock_gemini.side_effect = RuntimeError("Gemini error")
        mock_claude.side_effect = RuntimeError("Claude error")

        result = predict_success_rate(
            coverage_report=coverage_report,
            quality_metrics=quality_metrics,
            dataset_info=dataset_info,
            api_key="fake-key",
            anthropic_key="fake-key",
        )

    assert result.model_used == "heuristic"
    assert 0.0 <= result.predicted_success_rate <= 1.0


# ---------------------------------------------------------------------------
# Test 5: Malformed JSON with markdown
# ---------------------------------------------------------------------------

def test_malformed_json_response():
    """Markdown-wrapped JSON (```json ... ```) → parsed correctly."""
    text = '```json\n{"predicted_success_rate": 0.55, "confidence": 0.6, "grade": "C"}\n```'
    result = _parse_vlm_json(text)
    assert result["predicted_success_rate"] == 0.55
    assert result["grade"] == "C"


# ---------------------------------------------------------------------------
# Test 6: Partial JSON response
# ---------------------------------------------------------------------------

def test_partial_json_response():
    """JSON missing some fields → fills defaults without crashing."""
    partial = {"predicted_success_rate": 0.7, "reasoning": "Good dataset"}
    prediction = _build_prediction(partial, "gemini-2.5-flash", 100, 0.001)

    assert prediction.predicted_success_rate == 0.7
    assert prediction.confidence == 0.5  # default
    assert prediction.grade == "B"  # derived from 0.7
    assert prediction.reasoning == "Good dataset"
    assert prediction.similar_datasets == []
    assert prediction.key_strengths == []


# ---------------------------------------------------------------------------
# Test 7: Ground truth loading
# ---------------------------------------------------------------------------

def test_ground_truth_loading():
    """Verify ground_truth_condensed.json loads correctly."""
    # Reset cache
    import orbit.analyzer.vlm_predictor as vp
    vp._ground_truth_cache = None

    data = _load_ground_truth()
    assert len(data) == 82
    # Check first entry has required fields
    entry = data[0]
    assert "id" in entry
    assert "dataset" in entry
    assert "task_type" in entry
    assert "policy" in entry
    assert "success_rate" in entry
    assert "num_episodes" in entry
    assert "features_summary" in entry


# ---------------------------------------------------------------------------
# Test 8: Ground truth selection
# ---------------------------------------------------------------------------

def test_ground_truth_selection():
    """Given a pick_and_place task, selected entries are relevant."""
    data = _load_ground_truth()
    selected = _select_relevant_entries(data, task_type="pick_and_place", n=10)

    assert len(selected) == 10
    # At least some should be pick/place/grasp related
    relevant_keywords = {"pick", "place", "grasp", "lift", "can", "manipulation"}
    relevant_count = sum(
        1 for entry in selected
        if any(
            kw in entry.get("task_type", "").lower()
            or kw in entry.get("features_summary", "").lower()
            or kw in entry.get("dataset", "").lower()
            for kw in relevant_keywords
        )
    )
    assert relevant_count >= 3, f"Expected >=3 relevant entries, got {relevant_count}"


# ---------------------------------------------------------------------------
# Test 9: Grade mapping
# ---------------------------------------------------------------------------

def test_grade_mapping():
    """A for 0.85, B for 0.7, C for 0.55, D for 0.4, F for 0.2."""
    assert _grade_from_score(0.85) == "A"
    assert _grade_from_score(0.80) == "A"
    assert _grade_from_score(0.70) == "B"
    assert _grade_from_score(0.65) == "B"
    assert _grade_from_score(0.55) == "C"
    assert _grade_from_score(0.50) == "C"
    assert _grade_from_score(0.40) == "D"
    assert _grade_from_score(0.35) == "D"
    assert _grade_from_score(0.20) == "F"
    assert _grade_from_score(0.0) == "F"


# ---------------------------------------------------------------------------
# Test 10: Prompt construction
# ---------------------------------------------------------------------------

def test_prompt_construction(coverage_report, signal_diagnostics, quality_metrics, dataset_info):
    """Verify prompt includes coverage scores, diagnostics, and ground truth."""
    data = _load_ground_truth()
    refs = _select_relevant_entries(data, task_type="pick_and_place", n=10)

    prompt = _build_prompt(
        coverage_report=coverage_report,
        signal_diagnostics=signal_diagnostics,
        quality_metrics=quality_metrics,
        dataset_info=dataset_info,
        reference_entries=refs,
    )

    # Coverage scores present
    assert "position_diversity" in prompt
    assert "0.650" in prompt
    assert "action_diversity" in prompt

    # Signal diagnostics present
    assert "Signal Diagnostics" in prompt
    assert "WARNINGS" in prompt

    # Quality metrics
    assert "action_smoothness" in prompt

    # Reference entries
    assert "Reference Benchmarks" in prompt

    # Dataset info
    assert "pick up cups" in prompt


# ---------------------------------------------------------------------------
# Test 11: Cost tracking
# ---------------------------------------------------------------------------

def test_cost_tracking(coverage_report, dataset_info, valid_vlm_response):
    """Verify cost_estimate is populated and reasonable (< $0.05)."""
    with patch("orbit.analyzer.vlm_predictor._call_gemini") as mock_gemini:
        mock_gemini.return_value = (valid_vlm_response, 500, 0.003)

        result = predict_success_rate(
            coverage_report=coverage_report,
            dataset_info=dataset_info,
            api_key="fake-key",
        )

    assert result.cost_estimate > 0
    assert result.cost_estimate < 0.05


# ---------------------------------------------------------------------------
# Test 12: Empty diagnostics
# ---------------------------------------------------------------------------

def test_empty_diagnostics(coverage_report):
    """Pass minimal/empty diagnostics → doesn't crash."""
    empty_diag = SignalDiagnosticsReport(
        joint_diagnostics=[],
        gripper_analysis=GripperAnalysis(
            gripper_dim=None, open_ratio=0.0, close_ratio=0.0,
            transitions_per_episode=0.0, has_sufficient_grasps=False,
        ),
        episode_diagnostics=[],
        directional_balance=DirectionalBalance(
            per_joint_bias=[], overall_bias=0.5, dominant_direction=None,
        ),
        global_zero_velocity_ratio=0.0,
        dead_dimensions=[],
        episodes_to_remove=[],
        episodes_to_review=[],
    )

    with patch.dict("os.environ", {}, clear=True):
        env = {k: v for k, v in __import__("os").environ.items()
               if k not in ("GOOGLE_API_KEY", "ANTHROPIC_API_KEY")}
        with patch.dict("os.environ", env, clear=True):
            result = predict_success_rate(
                coverage_report=coverage_report,
                signal_diagnostics=empty_diag,
            )

    assert isinstance(result, VLMPrediction)
    assert result.model_used == "heuristic"


# ---------------------------------------------------------------------------
# Test 13: With frames
# ---------------------------------------------------------------------------

def test_with_frames(coverage_report, dataset_info, valid_vlm_response):
    """Pass sample_frames_base64 → verify they're included in the call."""
    frames = ["base64data1", "base64data2"]

    with patch("orbit.analyzer.vlm_predictor._call_gemini") as mock_gemini:
        mock_gemini.return_value = (valid_vlm_response, 500, 0.003)

        result = predict_success_rate(
            coverage_report=coverage_report,
            dataset_info=dataset_info,
            sample_frames_base64=frames,
            api_key="fake-key",
        )

    # Verify frames were passed to _call_gemini
    call_args = mock_gemini.call_args
    assert call_args is not None
    # frames is the 4th positional argument
    passed_frames = call_args[0][3] if len(call_args[0]) > 3 else call_args[1].get("frames")
    assert passed_frames == frames
    assert result.model_used == "gemini-2.5-flash"


# ---------------------------------------------------------------------------
# Test 14: Without frames
# ---------------------------------------------------------------------------

def test_without_frames(coverage_report, dataset_info, valid_vlm_response):
    """No frames → still works fine."""
    with patch("orbit.analyzer.vlm_predictor._call_gemini") as mock_gemini:
        mock_gemini.return_value = (valid_vlm_response, 500, 0.003)

        result = predict_success_rate(
            coverage_report=coverage_report,
            dataset_info=dataset_info,
            sample_frames_base64=None,
            api_key="fake-key",
        )

    assert result.model_used == "gemini-2.5-flash"
    assert result.predicted_success_rate == 0.62
