# filename: src/inference/model_loader.py
# purpose:  Load all inference models at FastAPI startup.
#           Checks local models/ first, falls back to HF Hub with 120s timeout.
#           Returns (models_dict, hashes_dict) — hashes stored separately from nn.Module.

from __future__ import annotations

import hashlib
import logging
from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout
from pathlib import Path
from typing import Any, Optional

import torch
import torch.nn as nn

import config as cfg
from src.models.base_classifier import load_model
from src.models.model_factory import get_model
from src.utils.helpers import download_model_from_hub

logger = logging.getLogger(__name__)

CNN_WEIGHT_FILES: dict[str, str] = {
    "resnet50": "resnet50_best.pt",
    "mobilenet": "mobilenet_best.pt",
}
YOLO_WEIGHT_PATH: Path = cfg.ARTIFACTS_DIR / "detection" / "yolov8_smartvision.pt"
HF_DOWNLOAD_TIMEOUT: int = 120


def load_all_models(device: str = "cpu") -> tuple[dict[str, Any], dict[str, str]]:
    """Load all inference models. Returns (models_dict, hashes_dict).

    models_dict keys: "resnet50", "mobilenet" (if FAST_MODE=False), "yolo"
    hashes_dict keys: same as models_dict (excludes "yolo")

    Gracefully skips any model that fails to load — caller checks which are present.
    FAST_MODE=True (default): ResNet50 + YOLO only.
    FAST_MODE=False: ResNet50 + MobileNet + YOLO.
    """
    models: dict[str, Any] = {}
    hashes: dict[str, str] = {}

    cnn_names = ["resnet50"]
    if not cfg.FAST_MODE:
        cnn_names.append("mobilenet")

    for name in cnn_names:
        model, model_hash = _load_cnn(name, device)
        if model is not None:
            models[name] = model
            hashes[name] = model_hash
            logger.info(f"Loaded CNN: {name} (hash={model_hash})")
        else:
            logger.warning(f"Skipping {name} — weight file not available.")

    yolo = _load_yolo()
    if yolo is not None:
        models["yolo"] = yolo
        logger.info("Loaded YOLO detector.")
    else:
        logger.warning("Skipping YOLO — weight file not available.")

    _warm_up(models, device)
    return models, hashes


def _load_cnn(name: str, device: str) -> tuple[Optional[nn.Module], str]:
    """Build architecture + load weights. Returns (model, hash_8chars) or (None, '')."""
    weight_path = _resolve_weight_path(name)
    if weight_path is None:
        return None, ""

    try:
        model_hash = _file_hash(weight_path)
        model = get_model(name, num_classes=cfg.NUM_CLASSES, pretrained=False)
        model = load_model(model, weight_path, device=device)
        model.eval()
        return model, model_hash
    except Exception as e:
        logger.error(f"Failed to load CNN {name}: {e}")
        return None, ""


def _file_hash(path: Path) -> str:
    """MD5 via 8MB chunks — avoids reading full file (100MB+) into memory."""
    hasher = hashlib.md5()
    with open(path, "rb") as f:
        while chunk := f.read(8 * 1024 * 1024):
            hasher.update(chunk)
    return hasher.hexdigest()[:8]


def _load_yolo() -> Optional[Any]:
    """Load YOLOv8n from local weight file.

    NOTE: YOLO .pt files require weights_only=False because they embed custom
    ultralytics classes (DetectionModel, etc.). This is safe here because the
    file was generated by our own training run in Section 6.
    """
    if not YOLO_WEIGHT_PATH.exists():
        logger.error(f"YOLO weight not found: {YOLO_WEIGHT_PATH}")
        return None
    try:
        from ultralytics import YOLO

        return YOLO(str(YOLO_WEIGHT_PATH))
    except Exception as e:
        logger.error(f"Failed to load YOLO: {e}")
        return None


def _warm_up(models: dict[str, Any], device: str) -> None:
    """3 forward passes per CNN to stabilize CPU JIT and kernel caching.
    YOLO excluded — ultralytics handles its own internal warmup.
    """
    cnn_models = {k: v for k, v in models.items() if k != "yolo"}
    if not cnn_models:
        return
    dummy = torch.randn(1, 3, cfg.IMAGE_SIZE, cfg.IMAGE_SIZE).to(device)
    with torch.no_grad():
        for name, model in cnn_models.items():
            for _ in range(3):
                _ = model(dummy)
    logger.info(f"Warmup complete (3 passes): {list(cnn_models.keys())}")


def _resolve_weight_path(name: str) -> Optional[Path]:
    """Check local models/ dir first, then try HF Hub download."""
    local_path = cfg.MODELS_DIR / CNN_WEIGHT_FILES[name]
    if local_path.exists():
        logger.info(f"Found local weights: {local_path}")
        return local_path

    logger.info(f"{local_path.name} not found locally. Attempting HF Hub download...")
    return _download_with_timeout(name)


def _download_with_timeout(name: str) -> Optional[Path]:
    """Download weight from HF Hub with HF_DOWNLOAD_TIMEOUT second limit.
    Returns local Path on success, None on timeout or any error.
    """
    if (
        not cfg.HF_TOKEN
        or not cfg.HF_REPO_ID
        or cfg.HF_REPO_ID == "your-username/smartvision-models"
    ):
        logger.warning(
            f"HF_TOKEN or HF_REPO_ID not configured. Cannot download {name}. "
            "Set HF_TOKEN and HF_REPO_ID in .env to enable automatic download."
        )
        return None

    filename = CNN_WEIGHT_FILES[name]

    def _do_download() -> Path:
        return download_model_from_hub(
            filename=filename,
            local_dir=cfg.MODELS_DIR,
            repo_id=cfg.HF_REPO_ID,
            token=cfg.HF_TOKEN,
        )

    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(_do_download)
        try:
            return future.result(timeout=HF_DOWNLOAD_TIMEOUT)
        except FuturesTimeout:
            logger.error(
                f"HF Hub download for {name} timed out after {HF_DOWNLOAD_TIMEOUT}s. "
                "Model will not be loaded."
            )
            return None
        except Exception as e:
            logger.error(f"HF Hub download failed for {name}: {e}")
            return None