"""Tests for the AI quality judge module.

Tests cover:
- Prompt formatting
- JSON response parsing (valid, markdown-wrapped, malformed)
- Graceful fallback when no API key / no google-genai / API error
- Result dataclass structure
"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import pytest

from orbit.analyzer.ai_quality_judge import (
    AIJudgeResult,
    _format_stats_for_judge,
    check_grade,
)


class TestFormatStats:
    def test_basic_formatting(self):
        result = _format_stats_for_judge(
            divergence=0.161, n_clipping=2, policy_fit=0.88,
            consistency=0.793, elevated_cvs=3,
        )
        assert "0.161" in result
        assert "2 joints" in result
        assert "0.88" in result
        assert "3/8" in result
        assert "0.793" in result

    def test_zero_values(self):
        result = _format_stats_for_judge(0.0, 0, 0.0, 0.0, 0)
        assert "0.000" in result
        assert "0 joints" in result
        assert "0/8" in result

    def test_custom_total_cvs(self):
        result = _format_stats_for_judge(0.1, 1, 0.9, 0.8, 2, total_cvs=10)
        assert "2/10" in result


class TestCheckGradeNoAPIKey:
    """Fallback when GOOGLE_API_KEY is not set."""

    @patch.dict("os.environ", {}, clear=True)
    def test_no_api_key_keeps_grade(self):
        result = check_grade(0.20, 3, 0.85, 0.70, 4)
        assert isinstance(result, AIJudgeResult)
        assert result.keep_grade is True
        assert result.error == "GOOGLE_API_KEY not set"
        assert result.model_used == "none"

    @patch.dict("os.environ", {}, clear=True)
    def test_returns_reasoning(self):
        result = check_grade(0.1, 0, 0.9, 0.8, 1)
        assert "No API key" in result.reasoning


class TestCheckGradeNoGenai:
    """Fallback when google-genai is not installed."""

    @patch.dict("os.environ", {"GOOGLE_API_KEY": "fake-key"})
    def test_import_error_keeps_grade(self):
        # Temporarily hide google.genai to trigger ImportError
        import sys
        saved = sys.modules.get("google.genai")
        saved_google = sys.modules.get("google")
        sys.modules["google.genai"] = None
        sys.modules["google"] = None
        try:
            result = check_grade(0.15, 2, 0.88, 0.79, 3)
            assert result.keep_grade is True
            assert result.error is not None
        finally:
            if saved is not None:
                sys.modules["google.genai"] = saved
            else:
                sys.modules.pop("google.genai", None)
            if saved_google is not None:
                sys.modules["google"] = saved_google
            else:
                sys.modules.pop("google", None)


def _make_mock_genai(response_text: str, side_effect=None):
    """Create a mock google module with genai.Client wired up.

    The function under test does ``from google import genai`` then
    ``genai.Client(api_key=...)``.  We need sys.modules["google"]
    to expose a ``.genai`` whose ``.Client(...)`` returns a mock
    client with the expected response.
    """
    mock_response = MagicMock()
    mock_response.text = response_text

    mock_client = MagicMock()
    if side_effect:
        mock_client.models.generate_content.side_effect = side_effect
    else:
        mock_client.models.generate_content.return_value = mock_response

    mock_genai = MagicMock()
    mock_genai.Client.return_value = mock_client

    mock_google = MagicMock()
    mock_google.genai = mock_genai

    return mock_google, mock_genai, mock_client


class TestCheckGradeMockedAPI:
    """Test with mocked Gemini API responses."""

    @patch.dict("os.environ", {"GOOGLE_API_KEY": "fake-key"})
    def test_keep_a_response(self):
        mock_google, mock_genai, _ = _make_mock_genai(
            '{"keep_A": true, "reasoning": "Low divergence, clean data."}'
        )
        with patch.dict("sys.modules", {"google": mock_google, "google.genai": mock_genai}):
            result = check_grade(0.09, 0, 0.92, 0.85, 1)
        assert result.keep_grade is True
        assert "Low divergence" in result.reasoning
        assert result.error is None
        assert result.model_used == "gemini-2.5-flash"

    @patch.dict("os.environ", {"GOOGLE_API_KEY": "fake-key"})
    def test_downgrade_response(self):
        mock_google, mock_genai, _ = _make_mock_genai(
            '{"keep_A": false, "reasoning": "Moderate divergence plus clipping."}'
        )
        with patch.dict("sys.modules", {"google": mock_google, "google.genai": mock_genai}):
            result = check_grade(0.16, 2, 0.88, 0.79, 3)
        assert result.keep_grade is False
        assert "divergence" in result.reasoning.lower()

    @patch.dict("os.environ", {"GOOGLE_API_KEY": "fake-key"})
    def test_markdown_wrapped_json(self):
        mock_google, mock_genai, _ = _make_mock_genai(
            '```json\n{"keep_A": false, "reasoning": "Issues found."}\n```'
        )
        with patch.dict("sys.modules", {"google": mock_google, "google.genai": mock_genai}):
            result = check_grade(0.18, 3, 0.85, 0.70, 4)
        assert result.keep_grade is False

    @patch.dict("os.environ", {"GOOGLE_API_KEY": "fake-key"})
    def test_malformed_json_keeps_grade(self):
        mock_google, mock_genai, _ = _make_mock_genai(
            "I think the grade should be B because the data has issues."
        )
        with patch.dict("sys.modules", {"google": mock_google, "google.genai": mock_genai}):
            result = check_grade(0.16, 2, 0.88, 0.79, 3)
        assert result.keep_grade is True
        assert result.error is not None
        assert "parse" in result.error.lower()

    @patch.dict("os.environ", {"GOOGLE_API_KEY": "fake-key"})
    def test_api_error_keeps_grade(self):
        mock_google, mock_genai, _ = _make_mock_genai(
            "", side_effect=RuntimeError("API down")
        )
        with patch.dict("sys.modules", {"google": mock_google, "google.genai": mock_genai}):
            result = check_grade(0.16, 2, 0.88, 0.79, 3)
        assert result.keep_grade is True
        assert result.error is not None
        assert "API down" in result.error

    @patch.dict("os.environ", {"GOOGLE_API_KEY": "fake-key"})
    def test_missing_fields_default_to_keep(self):
        mock_google, mock_genai, _ = _make_mock_genai('{"reasoning": "looks fine"}')
        with patch.dict("sys.modules", {"google": mock_google, "google.genai": mock_genai}):
            result = check_grade(0.10, 0, 0.90, 0.85, 1)
        assert result.keep_grade is True


class TestAIJudgeResult:
    def test_dataclass_fields(self):
        r = AIJudgeResult(keep_grade=True, reasoning="test", model_used="gemini")
        assert r.keep_grade is True
        assert r.reasoning == "test"
        assert r.model_used == "gemini"
        assert r.error is None

    def test_dataclass_with_error(self):
        r = AIJudgeResult(
            keep_grade=True, reasoning="fallback", model_used="none", error="no key"
        )
        assert r.error == "no key"
