from __future__ import annotations

from unittest.mock import MagicMock
from unittest.mock import patch

from amsdal_utils.models.enums import ModuleType

from amsdal_ml.mcp_server.server_model_explorer.services.types import ModelInfo


def test_get_llm_caches_and_reloads_on_class_change():
    from amsdal_ml.mcp_server.server_model_explorer.services import dependencies

    dependencies._LLM_INSTANCE = None
    dependencies._LLM_CLASS_PATH = None

    with patch(
        'amsdal_ml.mcp_server.server_model_explorer.services.dependencies.ml_config.mcp_ml_model_class',
        'class.one',
    ):
        fake_model_1 = MagicMock()
        fake_model_cls_1 = MagicMock(return_value=fake_model_1)

        with patch.object(dependencies, 'load_ml_model_class', return_value=fake_model_cls_1):
            llm_1 = dependencies._get_llm()
            llm_2 = dependencies._get_llm()

        assert llm_1 is llm_2
        fake_model_1.setup.assert_called_once()

    with patch(
        'amsdal_ml.mcp_server.server_model_explorer.services.dependencies.ml_config.mcp_ml_model_class',
        'class.two',
    ):
        fake_model_2 = MagicMock()
        fake_model_cls_2 = MagicMock(return_value=fake_model_2)

        with patch.object(dependencies, 'load_ml_model_class', return_value=fake_model_cls_2):
            llm_3 = dependencies._get_llm()

        assert llm_3 is fake_model_2
        fake_model_2.setup.assert_called_once()


def test_get_model_registry_caches_loaded_models():
    from amsdal_ml.mcp_server.server_model_explorer.services import dependencies
    from tests.fixtures.models.category import Category
    from tests.fixtures.models.user import User

    dependencies._MODEL_REGISTRY = None

    loaded_models = {
        User: ModuleType.USER,
        Category: ModuleType.CONTRIB,
    }

    with patch('amsdal_ml.utils.model_scanner.get_all_models', return_value=loaded_models):
        first = dependencies._get_model_registry()
        second = dependencies._get_model_registry()

    assert first is second
    assert first['User'] == ModelInfo(model=User, scope=ModuleType.USER)
    assert first['Category'] == ModelInfo(model=Category, scope=ModuleType.CONTRIB)
