"""
Model loading utilities for the inference server.

Provides a single ModelRegistry class that loads all trained artifacts
once at startup and exposes them to the API layer.

Usage:
    from scripts.serve_utils import ModelRegistry
    registry = ModelRegistry.load(cfg)
"""

from __future__ import annotations

import pickle
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd

from src.config import Config, MODELS_DIR, PROCESSED_DIR
from src.explainability.language_explainer import LanguageExplainer
from src.models.classical import LightGBMReranker
from src.models.naive import GlobalPopularityRecommender
from src.utils.logging import get_logger

logger = get_logger(__name__)

ITEM_META_COLS = [
    "item_avg_rating", "item_num_ratings", "item_popularity",
    "item_positive_rate", "price_norm", "category_idx", "brand_idx",
]


@dataclass
class ModelRegistry:
    """
    Holds all loaded model artifacts for the inference API.

    All attributes are populated by ModelRegistry.load(cfg).
    """
    naive_model: Optional[GlobalPopularityRecommender] = None
    lgbm_model: Optional[LightGBMReranker] = None
    deep_trainer: Optional[object] = None  # TwoTowerTrainer (lazy-imported)
    item_embeddings: Optional[np.ndarray] = None
    item_features: Optional[pd.DataFrame] = None
    user_features: Optional[pd.DataFrame] = None
    train_df: Optional[pd.DataFrame] = None
    feature_artifacts: dict = field(default_factory=dict)
    item_index: dict = field(default_factory=dict)
    language_explainer: Optional[LanguageExplainer] = None
    user_seen: dict[int, set[int]] = field(default_factory=dict)
    n_users: int = 0
    n_items: int = 0

    # Which models are available
    naive_available: bool = False
    classical_available: bool = False
    deep_available: bool = False

    @classmethod
    def load(cls, cfg: Config) -> "ModelRegistry":
        """
        Load all available model artifacts from disk.

        Missing artifacts are skipped with a warning (so partial deployments work).

        Args:
            cfg: Full Config object.

        Returns:
            Populated ModelRegistry.
        """
        registry = cls()

        # --- Shared data ---
        try:
            registry.item_features = pd.read_parquet(cfg.data.item_features_path)
            registry.n_items = int(registry.item_features["item_idx"].max() + 1)
        except Exception as e:
            logger.warning("Could not load item_features: %s", e)

        try:
            registry.user_features = pd.read_parquet(cfg.data.user_features_path)
        except Exception as e:
            logger.warning("Could not load user_features: %s", e)

        try:
            registry.train_df = pd.read_parquet(cfg.data.train_path)
            registry.n_users = int(registry.train_df["user_idx"].max() + 1)
            registry.user_seen = (
                registry.train_df.groupby("user_idx")["item_idx"].apply(set).to_dict()
            )
        except Exception as e:
            logger.warning("Could not load train_df: %s", e)

        try:
            fa_path = PROCESSED_DIR / "feature_artifacts.pkl"
            with open(fa_path, "rb") as f:
                registry.feature_artifacts = pickle.load(f)
        except Exception as e:
            logger.warning("Could not load feature_artifacts: %s", e)

        try:
            with open(cfg.deep.item_index_path, "rb") as f:
                registry.item_index = pickle.load(f)
        except Exception as e:
            logger.warning("Could not load item_index: %s", e)

        # --- Naive model ---
        naive_path = MODELS_DIR / "naive_baseline.pkl"
        if naive_path.exists():
            try:
                registry.naive_model = GlobalPopularityRecommender.load(naive_path)
                registry.naive_available = True
                logger.info("Naive model loaded.")
            except Exception as e:
                logger.warning("Failed to load naive model: %s", e)

        # --- Classical model ---
        lgbm_path = cfg.classical.model_path
        if lgbm_path.exists():
            try:
                registry.lgbm_model = LightGBMReranker.load(lgbm_path)
                registry.classical_available = True
                logger.info("LightGBM model loaded.")
            except Exception as e:
                logger.warning("Failed to load LightGBM model: %s", e)

        # --- Deep model ---
        if cfg.deep.model_path.exists() and cfg.deep.item_embeddings_path.exists():
            try:
                # Lazy import: torch must be imported AFTER LightGBM pickle load on Apple Silicon
                # to avoid a SIGSEGV caused by OpenMP/allocator conflicts.
                # Set single-threaded mode IMMEDIATELY after import to prevent OpenMP deadlock
                # when LightGBM's thread pool and PyTorch's OpenMP pool co-exist.
                logger.info("Deep model: importing torch…")
                import torch  # noqa: PLC0415
                torch.set_num_threads(1)
                torch.set_num_interop_threads(1)
                logger.info("Deep model: torch imported (single-thread mode), importing deep module…")
                from src.models.deep import TwoTowerModel, TwoTowerTrainer, build_item_feature_matrix  # noqa: PLC0415
                logger.info("Deep model: module imported, detecting device…")

                device = "cuda" if torch.cuda.is_available() else "cpu"
                logger.info("Deep model: device=%s", device)
                n_meta = 0
                if registry.item_features is not None:
                    logger.info("Deep model: building item feature matrix…")
                    item_feat_matrix = build_item_feature_matrix(
                        registry.item_features, ITEM_META_COLS
                    )
                    n_meta = item_feat_matrix.shape[1]
                    logger.info("Deep model: n_meta=%d", n_meta)
                else:
                    item_feat_matrix = None

                logger.info("Deep model: loading state dict from %s…", cfg.deep.model_path)
                state_dict = torch.load(cfg.deep.model_path, map_location=device, weights_only=True)
                logger.info("Deep model: state dict loaded (%d keys), creating model…", len(state_dict))
                model = TwoTowerModel(
                    n_users=registry.n_users,
                    n_items=registry.n_items,
                    embedding_dim=cfg.deep.embedding_dim,
                    output_dim=cfg.deep.embedding_dim,
                    hidden_dims=cfg.deep.hidden_dims,
                    dropout=0.0,          # no dropout at inference
                    n_meta_features=n_meta,
                )
                logger.info("Deep model: applying state dict…")
                model.load_state_dict(state_dict)
                del state_dict
                logger.info("Deep model: parameters loaded.")
                logger.info("Deep model: state dict loaded, creating trainer…")
                registry.deep_trainer = TwoTowerTrainer(
                    model=model, device=device, model_path=cfg.deep.model_path
                )
                logger.info("Deep model: loading embeddings…")
                registry.item_embeddings = np.load(cfg.deep.item_embeddings_path)
                registry.deep_available = True
                logger.info("Two-Tower model loaded. Item embeddings shape: %s",
                            registry.item_embeddings.shape)
            except Exception as e:
                logger.warning("Failed to load deep model: %s", e)

        # --- Language explainer (for deep model) ---
        if registry.train_df is not None and registry.item_features is not None:
            try:
                registry.language_explainer = LanguageExplainer(
                    registry.train_df,
                    registry.item_features,
                    positive_threshold=cfg.data.positive_rating_threshold,
                )
                logger.info("Language explainer initialised.")
            except Exception as e:
                logger.warning("Failed to initialise language explainer: %s", e)

        logger.info(
            "ModelRegistry ready — naive: %s | classical: %s | deep: %s",
            registry.naive_available,
            registry.classical_available,
            registry.deep_available,
        )
        return registry