from unittest.mock import MagicMock, patch

import pytest

from GramAddict.core.llm_provider import unload_ollama_models


def test_unload_ollama_models_sends_keep_alive_0():
    """
    Ensures that when unload_ollama_models is called, it correctly identifies
    local Ollama models from the config and sends a POST request with keep_alive: 0
    to unload them from VRAM.
    """
    mock_configs = MagicMock()
    mock_configs.args.ai_telepathic_model = "llama3.2:1b"
    mock_configs.args.ai_telepathic_url = "http://localhost:11434/api/generate"

    mock_configs.args.ai_fallback_model = "qwen2.5:latest"
    mock_configs.args.ai_fallback_url = "http://127.0.0.1:11434/api/generate"

    # Cloud model should NOT be unloaded
    mock_configs.args.ai_model = "openrouter/anthropic/claude"
    mock_configs.args.ai_model_url = "https://openrouter.ai/api/v1/chat/completions"

    with patch("GramAddict.core.llm_provider.requests.post") as mock_post:
        unload_ollama_models(mock_configs)

        # unload_ollama_models uses a background thread, so we must wait slightly or mock the threading.
        # But wait! We can just call the inner _unload directly, or wait a fraction of a second.
        import time

        time.sleep(0.1)

        # Expect 2 calls (for the 2 local models)
        assert mock_post.call_count == 2

        # Extract the JSON bodies of the calls
        called_json_args = [call.kwargs.get("json") for call in mock_post.call_args_list]

        # Verify keep_alive: 0 is present for both
        assert {"model": "llama3.2:1b", "keep_alive": 0} in called_json_args
        assert {"model": "qwen2.5:latest", "keep_alive": 0} in called_json_args

        # Verify cloud model was skipped
        assert not any(arg.get("model") == "openrouter/anthropic/claude" for arg in called_json_args)


def test_bot_flow_triggers_ollama_cleanup():
    """
    Ensures that the start_bot function triggers unload_ollama_models
    in its finally block when finishing or aborting.
    """
    from GramAddict.core.bot_flow import start_bot

    with (
        patch("GramAddict.core.bot_flow.Config") as mock_config_cls,
        patch("GramAddict.core.bot_flow.configure_logger"),
        patch("GramAddict.core.bot_flow.check_if_updated"),
        patch("GramAddict.core.benchmark_guard.check_model_benchmarks"),
        patch("GramAddict.core.llm_provider.log_openrouter_burn"),
        patch("GramAddict.core.llm_provider.prewarm_ollama_models"),
        patch("GramAddict.core.bot_flow.create_device") as mock_create_device,
        patch("GramAddict.core.session_state.SessionState.inside_working_hours", return_value=(True, 0)),
        patch("GramAddict.core.llm_provider.unload_ollama_models") as mock_unload,
    ):
        mock_configs = MagicMock()
        mock_config_cls.return_value = mock_configs

        # Simulate a crash inside the try block
        mock_device = MagicMock()
        mock_device.wake_up.side_effect = Exception("Simulate immediate crash")
        mock_create_device.return_value = mock_device

        with pytest.raises(Exception, match="Simulate immediate crash"):
            start_bot()

        # Verify the cleanup was STILL called even during a crash
        mock_unload.assert_called_once_with(mock_configs)
