"""Model loading / unloading for depth (DA3) and segmentation (SegEarth-OV3).

Model IDs and prompts come from config objects — nothing hardcoded.
"""

from __future__ import annotations

import gc
import logging
import tempfile
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn

from src.conf.hardware_conf import HardwareConfig
from src.conf.models_conf import ModelsConfig
from src.conf.seg_conf import SegConfig
from src.utils.profiler import profile_model, log_vram_snapshot

import src.nn  # noqa: F401  — registers vendored packages on sys.path

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Depth
# ---------------------------------------------------------------------------

def _resolve_cache_dir(models_conf: ModelsConfig) -> str | None:
    """Return absolute cache_dir for from_pretrained, or None for HF default."""
    if not models_conf.weights_dir:
        return None
    cache = Path(models_conf.weights_dir)
    if not cache.is_absolute():
        cache = Path(__file__).resolve().parents[2] / cache
    cache.mkdir(parents=True, exist_ok=True)
    return str(cache)


def load_depth_model(
    models_conf: ModelsConfig,
    hw_conf: HardwareConfig,
    device: torch.device,
) -> nn.Module:
    """Load depth estimation model from config.

    Args:
        models_conf: Model IDs from gin config.
        hw_conf: FP16 setting.
        device: Target CUDA device.

    Returns:
        Loaded depth model on device.
    """
    model_id = models_conf.depth_model_id
    cache_dir = _resolve_cache_dir(models_conf)
    logger.info("Loading depth model: %s (cache_dir=%s)", model_id, cache_dir)

    try:
        from depth_anything_3.api import DepthAnything3

        kwargs = {"cache_dir": cache_dir} if cache_dir else {}
        model = DepthAnything3.from_pretrained(f"depth-anything/{model_id}", **kwargs)
        model = model.to(device=device)
        model.eval()
        # DA3 forward is too complex for fvcore/thop (optional kwargs, nested models).
        # Profile params + VRAM only.
        profile_model(model, None, device, model_name=f"Depth ({model_id})")
        log_vram_snapshot("after depth load")
        return model
    except ImportError:
        logger.warning(
            "⚠️  DA3 not found, falling back to %s", models_conf.depth_fallback_id,
        )
        from transformers import AutoModelForDepthEstimation

        dtype = torch.float16 if hw_conf.use_fp16 else torch.float32
        model = AutoModelForDepthEstimation.from_pretrained(
            models_conf.depth_fallback_id, torch_dtype=dtype,
            cache_dir=cache_dir,
        ).to(device).eval()
        profile_model(model, (1, 3, 256, 256), device,
                       model_name=f"Depth ({models_conf.depth_fallback_id})")
        log_vram_snapshot("after depth load")
        return model


# ---------------------------------------------------------------------------
# CHMv2 (DINOv3 cross-view height map)
# ---------------------------------------------------------------------------

def load_chmv2_model(
    models_conf: ModelsConfig,
    hw_conf: HardwareConfig,
    device: torch.device,
) -> tuple[nn.Module, Any]:
    """Load CHMv2 depth model and processor.

    Returns:
        (model, processor) tuple.
    """
    cache_dir = _resolve_cache_dir(models_conf)
    # Prefer local weights if available.
    local_dir = Path(cache_dir) / "dinov3-chmv2" if cache_dir else None
    model_path = str(local_dir) if local_dir and local_dir.exists() else models_conf.chmv2_model_id

    logger.info("Loading CHMv2 model: %s", model_path)
    from transformers import CHMv2ForDepthEstimation, CHMv2ImageProcessor

    # CHMv2 (DINOv3 DPT) produces NaN in FP16 — always use FP32.
    processor = CHMv2ImageProcessor.from_pretrained(model_path)
    model = CHMv2ForDepthEstimation.from_pretrained(
        model_path, torch_dtype=torch.float32,
    ).to(device).eval()
    profile_model(model, (1, 3, 518, 518), device, model_name="CHMv2 (DINOv3)")
    log_vram_snapshot("after chmv2 load")
    return model, processor


# ---------------------------------------------------------------------------
# Segmentation
# ---------------------------------------------------------------------------

def load_segmentation_model(
    models_conf: ModelsConfig,
    hw_conf: HardwareConfig,
    seg_conf: SegConfig,
    device: torch.device,
) -> tuple[Any, dict[str, Any]]:
    """Load segmentation model from config.

    Args:
        models_conf: Model type and fallback IDs.
        hw_conf: FP16 setting.
        seg_conf: Text prompts for open-vocabulary segmentation.
        device: Target CUDA device.

    Returns:
        (model_or_pipeline, config_dict).
    """
    prompts = seg_conf.prompts
    logger.info("Loading segmentation (%s, %d classes) ...",
                models_conf.seg_model_type, len(prompts))

    if models_conf.seg_model_type == "segearth-ov3":
        try:
            from segearthov3_segmentor import SegEarthOV3Segmentation

            # Generate classname file from prompts.
            classname_file = tempfile.NamedTemporaryFile(
                mode="w", suffix=".txt", delete=False,
            )
            for prompt in prompts:
                classname_file.write(f"{prompt}\n")
            classname_file.close()

            seg_kwargs = dict(
                classname_path=classname_file.name,
                device=device,
                prob_thd=seg_conf.threshold,
                confidence_threshold=0.5,
                use_sem_seg=True,
                use_presence_score=True,
                use_transformer_decoder=True,
            )
            # Resolve bpe_path from weights_dir.
            cache_dir = _resolve_cache_dir(models_conf)
            if cache_dir:
                bpe = Path(cache_dir) / "bpe_simple_vocab_16e6.txt.gz"
                if bpe.exists():
                    seg_kwargs["bpe_path"] = str(bpe)
            if cache_dir:
                # Prefer SAM 3.1 weights, fall back to SAM 3.
                sam31_ckpt = Path(cache_dir) / "sam3.1" / "sam3.1_multiplex.pt"
                sam3_ckpt = Path(cache_dir) / "sam3" / "sam3.pt"
                if sam31_ckpt.exists():
                    seg_kwargs["checkpoint_path"] = str(sam31_ckpt)
                    logger.info("  Using SAM3.1 checkpoint: %s", sam31_ckpt)
                elif sam3_ckpt.exists():
                    seg_kwargs["checkpoint_path"] = str(sam3_ckpt)
                    logger.info("  Using SAM3 checkpoint: %s", sam3_ckpt)
            model = SegEarthOV3Segmentation(**seg_kwargs)
            logger.info("  🗺️  SegEarth-OV3 loaded. Prompts: %s", prompts)
            log_vram_snapshot("after segearth load")

            # Clean up temp file.
            Path(classname_file.name).unlink(missing_ok=True)

            return model, {
                "type": "segearth-ov3",
                "prompts": prompts,
                "num_classes": len(prompts),
            }
        except ImportError:
            logger.warning("⚠️  SegEarth-OV3 not found, falling back to SegFormer.")
        except Exception as exc:
            logger.warning("⚠️  SegEarth-OV3 load failed: %s, falling back.", exc)

    # SegFormer fallback.
    from transformers import (
        SegformerForSemanticSegmentation,
        SegformerImageProcessor,
    )

    cache_dir = _resolve_cache_dir(models_conf)
    model_id = models_conf.seg_fallback_id
    logger.info("Loading fallback: %s (cache_dir=%s)", model_id, cache_dir)
    dtype = torch.float16 if hw_conf.use_fp16 else torch.float32
    processor = SegformerImageProcessor.from_pretrained(model_id, cache_dir=cache_dir)
    model = SegformerForSemanticSegmentation.from_pretrained(
        model_id, torch_dtype=dtype, cache_dir=cache_dir,
    ).to(device).eval()
    profile_model(model, (1, 3, 640, 640), device,
                   model_name=f"Segmentation ({model_id})")
    log_vram_snapshot("after segformer load")
    return model, {
        "type": "segformer",
        "processor": processor,
        "num_classes": 150,
        "prompts": [],
    }


# ---------------------------------------------------------------------------
# Unload
# ---------------------------------------------------------------------------

def unload_model(model: Any | None) -> None:
    """Delete model and free GPU memory."""
    if model is None:
        return
    if hasattr(model, "model"):
        del model.model
    del model
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        log_vram_snapshot("after unload")
    logger.debug("🗑️  Model unloaded, CUDA cache cleared.")
