from __future__ import annotations

from amsdal_ml.mcp_server.server_model_explorer.services.types import ModelInfo
from amsdal_ml.mcp_server.server_model_explorer.services.types import ModelRegistry
from amsdal_ml.mcp_server.utils import load_ml_model_class
from amsdal_ml.ml_config import ml_config
from amsdal_ml.ml_models.models import MLModel

_MODEL_REGISTRY: ModelRegistry | None = None
_LLM_INSTANCE: MLModel | None = None
_LLM_CLASS_PATH: str | None = None


def _get_llm() -> MLModel:
    """Lazy loads and returns an instance of the configured ML model."""

    global _LLM_INSTANCE, _LLM_CLASS_PATH  # noqa: PLW0603

    class_path = ml_config.mcp_ml_model_class
    if _LLM_INSTANCE is None or _LLM_CLASS_PATH != class_path:
        model_cls = load_ml_model_class(class_path)
        _LLM_INSTANCE = model_cls()
        _LLM_INSTANCE.setup()
        _LLM_CLASS_PATH = class_path
    return _LLM_INSTANCE


def _get_model_registry() -> ModelRegistry:
    """Lazy loads and returns the model registry, which maps model class names to ModelInfo."""

    global _MODEL_REGISTRY  # noqa: PLW0603
    if _MODEL_REGISTRY is None:
        from amsdal_ml.utils.model_scanner import get_all_models

        loaded_models = get_all_models()
        _MODEL_REGISTRY = {
            model.__name__: ModelInfo(model=model, scope=scope) for model, scope in loaded_models.items()
        }
    return _MODEL_REGISTRY
