"""Model Registry for loading, caching, and serving trained model artifacts.

Supports eager loading at startup and lazy loading on first access.
Reports per-model health status.
"""

import json
import logging
from pathlib import Path
from typing import Any

import joblib

from app.core.exceptions import ModelNotLoadedError

logger = logging.getLogger(__name__)

# All registered model names and their artifact subdirectory names
REGISTERED_MODELS: list[str] = [
    "lo_tagger",
    "bloom_classifier",
    "mastery_model",
    "risk_model",
    "answer_scorer",
    "recommender",
]


class ModelRegistry:
    """Loads and manages trained model artifacts.

    Supports eager loading at startup and lazy loading on first access.
    Reports per-model health status.
    """

    def __init__(self, artifact_dir: str | Path) -> None:
        self._artifact_dir = Path(artifact_dir)
        self._models: dict[str, dict[str, Any]] = {}
        self._status: dict[str, str] = {}  # "loaded" | "not_loaded" | "error"
        self._metadata: dict[str, dict] = {}  # metrics.json content per model

        # Initialize all registered models as not_loaded
        for model_name in REGISTERED_MODELS:
            self._status[model_name] = "not_loaded"

    def load_all(self) -> None:
        """Eagerly load all model artifacts from artifact_dir subdirectories."""
        for model_name in REGISTERED_MODELS:
            self._load_model(model_name)

    def _load_model(self, model_name: str) -> None:
        """Load a single model's artifacts from its subdirectory.

        On failure, sets status to 'error' and logs the exception without crashing.
        """
        model_dir = self._artifact_dir / model_name

        if not model_dir.exists():
            logger.warning(
                "Model directory not found for '%s': %s", model_name, model_dir
            )
            self._status[model_name] = "not_loaded"
            return

        try:
            model_data: dict[str, Any] = {}

            # Required: model.joblib
            model_path = model_dir / "model.joblib"
            if not model_path.exists():
                logger.warning(
                    "model.joblib not found for '%s' at %s", model_name, model_path
                )
                self._status[model_name] = "not_loaded"
                return

            model_data["model"] = joblib.load(model_path)

            # Optional: vectorizer.joblib
            vectorizer_path = model_dir / "vectorizer.joblib"
            if vectorizer_path.exists():
                model_data["vectorizer"] = joblib.load(vectorizer_path)

            # Optional: label_encoder.joblib
            label_encoder_path = model_dir / "label_encoder.joblib"
            if label_encoder_path.exists():
                model_data["label_encoder"] = joblib.load(label_encoder_path)

            # Optional: feature_columns.json
            feature_columns_path = model_dir / "feature_columns.json"
            if feature_columns_path.exists():
                with open(feature_columns_path, "r", encoding="utf-8") as f:
                    model_data["feature_columns"] = json.load(f)

            # Optional: metrics.json (stored as metadata)
            metrics_path = model_dir / "metrics.json"
            if metrics_path.exists():
                with open(metrics_path, "r", encoding="utf-8") as f:
                    self._metadata[model_name] = json.load(f)

            self._models[model_name] = model_data
            self._status[model_name] = "loaded"
            logger.info("Model '%s' loaded successfully from %s", model_name, model_dir)

        except Exception:
            logger.exception("Failed to load model '%s' from %s", model_name, model_dir)
            self._status[model_name] = "error"

    def get_model(self, model_name: str) -> dict[str, Any]:
        """Retrieve a loaded model by name. Triggers lazy load if not yet loaded.

        Returns a dict containing 'model' and optional 'vectorizer',
        'label_encoder', 'feature_columns' keys.

        Raises ModelNotLoadedError if artifact is missing/corrupted after load attempt.
        """
        if model_name not in self._status:
            raise ModelNotLoadedError(
                f"Model '{model_name}' is not registered in the registry."
            )

        # Lazy loading: attempt to load if not yet loaded
        if self._status[model_name] == "not_loaded":
            self._load_model(model_name)

        if self._status[model_name] != "loaded":
            raise ModelNotLoadedError(
                f"Model '{model_name}' is not available (status: {self._status[model_name]})."
            )

        return self._models[model_name]

    def get_status(self, model_name: str) -> str:
        """Return 'loaded', 'not_loaded', or 'error' for a given model."""
        return self._status.get(model_name, "not_loaded")

    def get_metadata(self, model_name: str) -> dict | None:
        """Return metrics.json content for a model, or None if not available."""
        return self._metadata.get(model_name)

    def get_all_status(self) -> dict[str, dict]:
        """Return status + metadata for all registered models."""
        result: dict[str, dict] = {}
        for model_name in REGISTERED_MODELS:
            entry: dict[str, Any] = {
                "status": self._status.get(model_name, "not_loaded"),
            }
            metadata = self._metadata.get(model_name)
            if metadata:
                entry["metadata"] = metadata
            result[model_name] = entry
        return result

    def is_loaded(self, model_name: str) -> bool:
        """Check if a model is loaded and ready for inference."""
        return self._status.get(model_name) == "loaded"