"""
ML model loader.

Production rule: if a model fails to load at startup, the application
MUST refuse to start, not start in a half-working state and fail on the
first real request. A health check that says "I'm up" while secretly
missing the yield model is worse than a crash — it's a silent SLA
violation the client discovers from a support ticket, not from monitoring.

All 6 models are loaded once, here, and held in memory for the lifetime
of the worker process. They are never reloaded per-request.
"""

from pathlib import Path
from typing import Any

import joblib

from app.core.config import get_settings
from app.core.exceptions import ModelNotLoadedError
from app.core.logging_config import get_logger

settings = get_settings()
logger = get_logger(__name__)

REQUIRED_MODEL_FILES = {
    "sarimax_climate": "sarimax_climate.pkl",
    "xgb_feasibility": "xgb_feasibility.pkl",
    "xgb_yield": "xgb_yield.pkl",
    "dt_trend": "dt_trend.pkl",
    "rf_recommendation": "rf_recommendation.pkl",
    "sarimax_market": "sarimax_market.pkl",
}


class ModelRegistry:
    """Holds every loaded model in memory and exposes typed accessors."""

    def __init__(self) -> None:
        self._models: dict[str, Any] = {}
        self._loaded = False
        self._model_dir = Path(settings.MODEL_DIR)

    def load_all(self) -> None:
        logger.info(f"Loading models from {self._model_dir} ...")
        missing = []
        failed = []

        for key, filename in REQUIRED_MODEL_FILES.items():
            path = self._model_dir / filename
            if not path.exists():
                missing.append(str(path))
                continue
            try:
                self._models[key] = joblib.load(path)
                logger.info(f"Loaded model '{key}' from {path}")
            except Exception as exc:
                failed.append((str(path), str(exc)))

        if missing:
            raise ModelNotLoadedError(
                message="Required model files are missing. The service cannot start.",
                detail=f"Missing files: {missing}",
            )
        if failed:
            raise ModelNotLoadedError(
                message="One or more model files failed to load. The service cannot start.",
                detail=f"Failed to load: {failed}",
            )

        self._validate_loaded_models()
        self._loaded = True
        logger.info(f"All {len(self._models)} models loaded successfully")

    def _validate_loaded_models(self) -> None:
        """
        Sanity-check the structure of each loaded artifact. Catches the case
        where a .pkl file loads without raising an exception but doesn't
        contain what the inference code expects (e.g. wrong dict shape from
        a bad training run accidentally copied to the model directory).
        """
        bundle_models = ["xgb_feasibility", "xgb_yield", "dt_trend", "rf_recommendation"]
        for key in bundle_models:
            obj = self._models.get(key)
            if not isinstance(obj, dict) or "model" not in obj:
                raise ModelNotLoadedError(
                    message="A model file is present but has an invalid format.",
                    detail=f"'{key}' expected dict with 'model' key, got {type(obj)}",
                )

        for key in ("sarimax_climate", "sarimax_market"):
            obj = self._models.get(key)
            if not isinstance(obj, dict):
                raise ModelNotLoadedError(
                    message="A model file is present but has an invalid format.",
                    detail=f"'{key}' expected dict of {{key: fitted_model}}, got {type(obj)}",
                )

    def _ensure_loaded(self) -> None:
        if not self._loaded:
            raise ModelNotLoadedError(
                message="Models are not yet loaded. Service is starting up.",
                detail="ModelRegistry accessed before load_all() completed",
            )

    def get(self, key: str) -> Any:
        self._ensure_loaded()
        if key not in self._models:
            raise ModelNotLoadedError(message=f"Model '{key}' was not loaded")
        return self._models[key]

    @property
    def is_ready(self) -> bool:
        return self._loaded

    @property
    def loaded_model_keys(self) -> list[str]:
        return list(self._models.keys())


# Single shared instance — populated once during app startup (see app/main.py)
# and once per Celery worker process (see app/workers/celery_app.py)
model_registry = ModelRegistry()