from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock
from unittest.mock import patch

import pytest
from amsdal_utils.models.enums import ModuleType

from amsdal_ml.mcp_server.server_model_explorer.services.types import ModelInfo
from tests.fixtures.models.book import Book
from tests.fixtures.models.catalog_item import CatalogItem
from tests.fixtures.models.category import Category
from tests.fixtures.models.product import Product
from tests.fixtures.models.user import User


@pytest.fixture
def mock_registry(async_test_amsdal_manager: Any):  # noqa: ARG001
    with (
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.discovery._get_model_registry') as mock_discovery,
        patch('amsdal_ml.mcp_server.server_model_explorer.tools.crud._get_model_registry') as mock_crud,
    ):
        registry = {'User': ModelInfo(model=User, scope=ModuleType.USER)}
        mock_discovery.return_value = registry
        mock_crud.return_value = registry
        yield mock_discovery


@pytest.fixture
def mock_registry_with_relationships(async_test_amsdal_manager: Any):  # noqa: ARG001
    with patch('amsdal_ml.mcp_server.server_model_explorer.tools.discovery._get_model_registry') as mock:
        mock.return_value = {
            'Category': ModelInfo(model=Category, scope=ModuleType.USER),
            'CatalogItem': ModelInfo(model=CatalogItem, scope=ModuleType.USER),
            'Book': ModelInfo(model=Book, scope=ModuleType.USER),
            'Product': ModelInfo(model=Product, scope=ModuleType.CONTRIB),
        }
        yield mock


@pytest.fixture
def mock_llm():
    with patch('amsdal_ml.mcp_server.server_model_explorer.tools.crud._get_llm') as mock:
        llm_instance = MagicMock()
        mock.return_value = llm_instance
        yield llm_instance
