"""Tests for structured Gemini vision quality assessment.

Tests prompt construction, JSON parsing, and grade adjustment computation.
All tests mock the Gemini API — no API key needed.
"""

import json
from unittest.mock import MagicMock, patch

import pytest

from orbit.analyzer.vision_quality import (
    VisionQualityResult,
    _compute_adjustments,
    _extract_frame_indices,
    _has_image_observations,
    _parse_response,
    _select_episodes,
    assess_vision_quality,
)


class TestHasImageObservations:
    def test_no_images(self):
        info = MagicMock()
        info.observation_keys = ["state", "action"]
        info.camera_names = None
        info.features = {}
        assert not _has_image_observations(info)

    def test_has_image_key(self):
        info = MagicMock()
        info.observation_keys = ["observation.image", "state"]
        info.camera_names = None
        info.features = {}
        assert _has_image_observations(info)

    def test_has_camera_names(self):
        info = MagicMock()
        info.observation_keys = []
        info.camera_names = ["top", "wrist"]
        info.features = {}
        assert _has_image_observations(info)

    def test_has_image_features(self):
        info = MagicMock()
        info.observation_keys = []
        info.camera_names = None
        info.features = {"observation.images.top": "Image(...)"}
        assert _has_image_observations(info)


class TestSelectEpisodes:
    def test_few_episodes(self):
        assert _select_episodes(None, 2) == [0, 1]

    def test_no_qualities_even_spacing(self):
        result = _select_episodes(None, 100)
        assert len(result) == 3
        assert result == [0, 50, 99]

    def test_with_qualities(self):
        qualities = {0: 0.1, 1: 0.5, 2: 0.9, 3: 0.3, 4: 0.7}
        result = _select_episodes(qualities, 5)
        assert len(result) == 3
        # Best should be 2 (0.9), worst should be 0 (0.1)
        assert 2 in result
        assert 0 in result

    def test_empty(self):
        assert _select_episodes(None, 0) == []


class TestExtractFrameIndices:
    def test_normal_episode(self):
        indices = _extract_frame_indices(100)
        assert len(indices) == 3
        assert indices[0] == 10  # 10%
        assert indices[1] == 50  # 50%
        assert indices[2] == 90  # 90%

    def test_short_episode(self):
        indices = _extract_frame_indices(2)
        assert indices == [0, 1]

    def test_empty(self):
        assert _extract_frame_indices(0) == []


class TestParseResponse:
    def test_valid_json(self):
        text = '{"episodes": [], "consistent_task": {"answer": true}}'
        result = _parse_response(text)
        assert result["consistent_task"]["answer"] is True

    def test_markdown_wrapped(self):
        text = '```json\n{"episodes": []}\n```'
        result = _parse_response(text)
        assert "episodes" in result

    def test_invalid_json(self):
        result = _parse_response("not json at all")
        assert result == {}


class TestComputeAdjustments:
    def test_all_good(self):
        parsed = {
            "episodes": [
                {"task_visible": {"answer": True}, "task_complete": {"answer": True},
                 "visual_quality": {"answer": True}, "anomaly": {"answer": False}},
            ],
            "consistent_task": {"answer": True},
        }
        adj = _compute_adjustments(parsed)
        assert len(adj) == 0

    def test_invisible_task(self):
        parsed = {
            "episodes": [
                {"task_visible": {"answer": False}, "task_complete": {"answer": True},
                 "visual_quality": {"answer": True}, "anomaly": {"answer": False}},
            ],
            "consistent_task": {"answer": True},
        }
        adj = _compute_adjustments(parsed)
        assert "camera_misconfigured" in adj
        assert adj["camera_misconfigured"] == -5

    def test_incomplete_tasks(self):
        parsed = {
            "episodes": [
                {"task_visible": {"answer": True}, "task_complete": {"answer": False},
                 "visual_quality": {"answer": True}, "anomaly": {"answer": False}},
                {"task_visible": {"answer": True}, "task_complete": {"answer": False},
                 "visual_quality": {"answer": True}, "anomaly": {"answer": False}},
            ],
            "consistent_task": {"answer": True},
        }
        adj = _compute_adjustments(parsed)
        assert "incomplete_tasks" in adj
        assert adj["incomplete_tasks"] == -10

    def test_inconsistent_tasks(self):
        parsed = {
            "episodes": [
                {"task_visible": {"answer": True}, "task_complete": {"answer": True},
                 "visual_quality": {"answer": True}, "anomaly": {"answer": False}},
            ],
            "consistent_task": {"answer": False},
        }
        adj = _compute_adjustments(parsed)
        assert "inconsistent_tasks" in adj
        assert adj["inconsistent_tasks"] == -8

    def test_poor_visual_quality(self):
        parsed = {
            "episodes": [
                {"task_visible": {"answer": True}, "task_complete": {"answer": True},
                 "visual_quality": {"answer": False}, "anomaly": {"answer": False}},
            ],
            "consistent_task": {"answer": True},
        }
        adj = _compute_adjustments(parsed)
        assert "poor_visual_quality" in adj
        assert adj["poor_visual_quality"] == -3

    def test_empty_episodes(self):
        adj = _compute_adjustments({"episodes": []})
        assert len(adj) == 0


class TestAssessVisionQuality:
    def test_state_only_dataset_skips(self):
        info = MagicMock()
        info.observation_keys = ["state"]
        info.camera_names = None
        info.features = {}
        result = assess_vision_quality(info)
        assert result.total_penalty == 0
        assert result.error == "No image observations in dataset"

    def test_no_frame_loader_skips(self):
        info = MagicMock()
        info.observation_keys = ["observation.image"]
        info.camera_names = None
        info.features = {}
        result = assess_vision_quality(info, frame_loader=None)
        assert result.total_penalty == 0
        assert result.error == "No frame loader provided"

    @patch.dict("os.environ", {}, clear=True)
    def test_no_api_key_skips(self):
        info = MagicMock()
        info.observation_keys = ["observation.image"]
        info.camera_names = None
        info.features = {}
        result = assess_vision_quality(info, frame_loader=lambda ep, f: None)
        assert result.total_penalty == 0
        assert "API_KEY" in (result.error or "")
