from unittest.mock import MagicMock, patch

from GramAddict.core.llm_provider import query_llm


def test_query_llm_hallucination_recovery():
    # Test that when the primary model hallucinates non-JSON, it triggers fallback
    with patch("requests.post") as mock_post:
        # 1st call: Primary fails entirely (e.g., Timeout or strange error)
        mock_response_1 = MagicMock()
        mock_response_1.status_code = 500
        mock_response_1.raise_for_status.side_effect = Exception("500 Server Error")

        # 2nd call: Fallback works and returns valid JSON
        mock_response_2 = MagicMock()
        mock_response_2.status_code = 200
        mock_response_2.raise_for_status.return_value = None
        mock_response_2.json.return_value = {"choices": [{"message": {"content": '{"test": "success"}'}}]}

        mock_post.side_effect = [mock_response_1, mock_response_2]

        # Attempt a query with a primary model
        res = query_llm(
            url="http://fake.api/v1/chat/completions",
            model="primary-model",
            prompt="Hello",
            format_json=True,
            fallback_model="fallback-model",
            fallback_url="http://fake.api/v1/chat/completions",
        )

        assert res is not None
        assert "response" in res
        assert res["response"] == '{"test": "success"}'
        assert mock_post.call_count == 2


def test_query_llm_double_hallucination_safe_return():
    # Test that when both models hallucinate, we return None gracefully
    with patch("requests.post") as mock_post:
        # Both models fail
        mock_response = MagicMock()
        mock_response.status_code = 500
        mock_response.raise_for_status.side_effect = Exception("500 Server Error")

        mock_post.side_effect = [mock_response, mock_response]

        res = query_llm(
            url="http://fake.api/v1/chat/completions", model="primary-model", prompt="Hello", format_json=True
        )

        assert res is None
