"""Model discovery, loading, and GPU residency for NormalCrafter.

NormalCrafter uses its own ComfyUI model category: ``normalcrafter``.
All path discovery goes through ComfyUI's ``folder_paths`` registry, so paths
registered through ``extra_model_paths.yaml`` are visible automatically.
"""

from __future__ import annotations

from dataclasses import dataclass
import gc
import logging
import os

import folder_paths
import torch
from diffusers import AutoencoderKLTemporalDecoder, EulerDiscreteScheduler
from huggingface_hub import snapshot_download
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection

from .unet import NormalCrafterUNet

logger = logging.getLogger(__name__)

MODEL_CATEGORY = "normalcrafter"
NORMALCRAFTER_REPO = "Yanrui95/NormalCrafter"
SVD_REPO = "stabilityai/stable-video-diffusion-img2vid-xt"

HF_PREFIX = "HF: "
LOCAL_PREFIX = "local: "


def _register_model_category() -> None:
    """Ensure ComfyUI knows the default ``models/normalcrafter`` location.

    If ``extra_model_paths.yaml`` already registered one or more ``normalcrafter``
    paths, they stay in ComfyUI's ordering and the standard ComfyUI model path is
    simply added as another search location.
    """

    default_path = os.path.join(folder_paths.models_dir, MODEL_CATEGORY)

    try:
        existing = folder_paths.get_folder_paths(MODEL_CATEGORY)
    except KeyError:
        existing = []

    if default_path in existing:
        return

    # Current ComfyUI supports ``is_default``. The fallback keeps compatibility
    # with older versions whose function accepted only two arguments.
    try:
        folder_paths.add_model_folder_path(
            MODEL_CATEGORY,
            default_path,
            is_default=not existing,
        )
    except TypeError:
        folder_paths.add_model_folder_path(MODEL_CATEGORY, default_path)


_register_model_category()


def _model_roots() -> list[str]:
    """Return all NormalCrafter roots currently registered in ComfyUI."""

    roots = folder_paths.get_folder_paths(MODEL_CATEGORY)
    # Preserve ComfyUI ordering while removing duplicate spellings of a path.
    seen: set[str] = set()
    result: list[str] = []
    for root in roots:
        normalized = os.path.normcase(os.path.abspath(root))
        if normalized in seen:
            continue
        seen.add(normalized)
        result.append(os.path.abspath(root))
    return result


def _is_within_registered_root(path: str) -> bool:
    candidate = os.path.normcase(os.path.abspath(path))
    for root in _model_roots():
        registered = os.path.normcase(os.path.abspath(root))
        try:
            if os.path.commonpath((candidate, registered)) == registered:
                return True
        except ValueError:
            # Different Windows drives cannot share a common path.
            continue
    return False


def _local_model_dirs(required_subfolders: tuple[str, ...]) -> list[str]:
    """Find compatible model directories across every ComfyUI-registered root."""

    results: list[str] = []

    for root in _model_roots():
        if not os.path.isdir(root):
            continue

        try:
            names = sorted(os.listdir(root), key=str.lower)
        except OSError:
            continue

        for name in names:
            path = os.path.join(root, name)
            if not os.path.isdir(path):
                continue
            if all(
                os.path.isdir(os.path.join(path, subfolder))
                for subfolder in required_subfolders
            ):
                results.append(f"{LOCAL_PREFIX}{path}")

    return results


def normalcrafter_sources() -> list[str]:
    """Dropdown values suitable for NormalCrafter UNet + VAE loading."""

    return [
        f"{HF_PREFIX}{NORMALCRAFTER_REPO}",
        *_local_model_dirs(("unet", "vae")),
    ]


def base_model_sources() -> list[str]:
    """Dropdown values suitable for SVD CLIP + scheduler loading."""

    return [
        f"{HF_PREFIX}{SVD_REPO}",
        *_local_model_dirs(("image_encoder", "feature_extractor", "scheduler")),
    ]


def _existing_repo_dir(repo_id: str) -> str | None:
    """Return an already-downloaded repo from any registered model root."""

    folder_name = repo_id.replace("/", "--")
    for root in _model_roots():
        candidate = os.path.join(root, folder_name)
        if os.path.isdir(candidate):
            return candidate
    return None


def _download_root() -> str:
    """Choose the first writable NormalCrafter root in ComfyUI's own order."""

    errors: list[str] = []
    for root in _model_roots():
        try:
            os.makedirs(root, exist_ok=True)
            if os.access(root, os.W_OK):
                return root
            errors.append(f"{root} (not writable)")
        except OSError as exc:
            errors.append(f"{root} ({exc})")

    details = "; ".join(errors) or "no paths registered"
    raise RuntimeError(
        "No writable ComfyUI 'normalcrafter' model path is available: " + details
    )


def _download_repo(repo_id: str, required_subfolders: tuple[str, ...]) -> str:
    """Download only required components into a registered NormalCrafter root."""

    existing = _existing_repo_dir(repo_id)
    if existing is not None:
        return existing

    target = os.path.join(_download_root(), repo_id.replace("/", "--"))
    allow_patterns = [f"{subfolder}/**" for subfolder in required_subfolders]

    snapshot_download(
        repo_id=repo_id,
        local_dir=target,
        allow_patterns=allow_patterns,
    )
    return target


def _resolve_source(source: str, required_subfolders: tuple[str, ...]) -> str:
    """Resolve a dropdown value to a validated local model directory."""

    if source.startswith(HF_PREFIX):
        path = _download_repo(source[len(HF_PREFIX) :], required_subfolders)
    elif source.startswith(LOCAL_PREFIX):
        path = source[len(LOCAL_PREFIX) :]
        if not _is_within_registered_root(path):
            raise ValueError(
                "Selected model path is not inside a ComfyUI-registered "
                f"'{MODEL_CATEGORY}' directory: {path}"
            )
    else:
        raise ValueError(f"Unsupported model source: {source}")

    if not os.path.isdir(path):
        raise FileNotFoundError(f"NormalCrafter model folder not found: {path}")

    missing = [
        subfolder
        for subfolder in required_subfolders
        if not os.path.isdir(os.path.join(path, subfolder))
    ]
    if missing:
        raise FileNotFoundError(
            f"Model folder {path} is missing required subfolders: {', '.join(missing)}"
        )

    return path


def _clear_cuda_cache() -> None:
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def _auto_dtype(device: torch.device) -> torch.dtype:
    return torch.float16 if device.type == "cuda" else torch.float32


@dataclass
class NormalCrafterModels:
    """Own the heavy modules and keep GPU residency explicit."""

    unet: NormalCrafterUNet
    vae: AutoencoderKLTemporalDecoder
    image_encoder: CLIPVisionModelWithProjection
    feature_extractor: CLIPImageProcessor
    scheduler: EulerDiscreteScheduler
    dtype: torch.dtype

    def _modules(self) -> dict[str, torch.nn.Module]:
        return {
            "image_encoder": self.image_encoder,
            "vae": self.vae,
            "unet": self.unet,
        }

    def activate(self, name: str, device: torch.device) -> torch.nn.Module:
        """Keep only one heavy component on the inference device."""

        modules = self._modules()
        if name not in modules:
            raise ValueError(f"Unknown NormalCrafter component: {name}")

        cpu = torch.device("cpu")
        for module_name, module in modules.items():
            target = device if module_name == name else cpu
            try:
                current = next(module.parameters()).device
            except StopIteration:
                current = cpu

            if current != target:
                module.to(target)

        _clear_cuda_cache()
        return modules[name]

    def offload_all(self) -> None:
        """Return every heavy module to CPU and release unused CUDA blocks."""

        for module in self._modules().values():
            module.to("cpu")
        _clear_cuda_cache()


def _enable_memory_efficient_attention(unet: NormalCrafterUNet) -> None:
    try:
        unet.enable_xformers_memory_efficient_attention()
        logger.info("NormalCrafter: xFormers attention enabled")
    except Exception as exc:
        logger.info("NormalCrafter: using Diffusers default attention (%s)", exc)


def load_models(
    *,
    device: torch.device,
    normalcrafter_source: str,
    base_model_source: str,
) -> NormalCrafterModels:
    """Load the selected model components on CPU.

    UNet and CLIP use FP16 on CUDA systems and FP32 on CPU. The temporal VAE is
    intentionally loaded in FP32, matching its force-upcast inference behavior.
    Inference moves CLIP, VAE and UNet to the GPU one stage at a time.
    """

    normalcrafter_path = _resolve_source(
        normalcrafter_source,
        ("unet", "vae"),
    )
    base_model_path = _resolve_source(
        base_model_source,
        ("image_encoder", "feature_extractor", "scheduler"),
    )

    dtype = _auto_dtype(device)
    common = {
        "torch_dtype": dtype,
        "low_cpu_mem_usage": True,
        "local_files_only": True,
    }

    unet = NormalCrafterUNet.from_pretrained(
        normalcrafter_path,
        subfolder="unet",
        **common,
    )

    # The temporal VAE declares force_upcast and is evaluated in FP32 in the
    # reference NormalCrafter pipeline. Load it in FP32 from the start instead of
    # repeatedly calling vae.to(dtype=...) during inference. Recent Diffusers
    # versions explicitly warn against those runtime dtype casts because selected
    # VAE modules may need to remain in FP32.
    vae = AutoencoderKLTemporalDecoder.from_pretrained(
        normalcrafter_path,
        subfolder="vae",
        torch_dtype=torch.float32,
        low_cpu_mem_usage=True,
        local_files_only=True,
    )

    encoder_kwargs = dict(common)
    if dtype == torch.float16:
        encoder_kwargs["variant"] = "fp16"

    try:
        image_encoder = CLIPVisionModelWithProjection.from_pretrained(
            base_model_path,
            subfolder="image_encoder",
            **encoder_kwargs,
        )
    except (OSError, ValueError, TypeError):
        encoder_kwargs.pop("variant", None)
        image_encoder = CLIPVisionModelWithProjection.from_pretrained(
            base_model_path,
            subfolder="image_encoder",
            **encoder_kwargs,
        )

    feature_extractor = CLIPImageProcessor.from_pretrained(
        base_model_path,
        subfolder="feature_extractor",
        local_files_only=True,
    )
    scheduler = EulerDiscreteScheduler.from_pretrained(
        base_model_path,
        subfolder="scheduler",
        local_files_only=True,
    )

    unet.eval()
    vae.eval()
    image_encoder.eval()
    _enable_memory_efficient_attention(unet)

    models = NormalCrafterModels(
        unet=unet,
        vae=vae,
        image_encoder=image_encoder,
        feature_extractor=feature_extractor,
        scheduler=scheduler,
        dtype=dtype,
    )
    models.offload_all()
    return models