"""Runtime wrapper for local single-GPU DiffiT inference inside ComfyUI."""

from __future__ import annotations

import importlib
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import torch

from .utils import (
    LOGGER,
    default_cfg_scale_for_image_size,
    get_hf_cache_dir,
    get_preferred_diffit_download_dir,
    infer_image_size_from_state_dict,
    latent_size_for_image,
    load_state_dict_file,
    make_generator,
    pick_device,
    resolve_diffit_checkpoint_path,
    to_comfy_image,
    validate_image_size,
)
from .vendor.diffit_upstream import create_diffusion, diffusion_defaults
from .vendor.diffit_upstream.script_util import NUM_CLASSES


@dataclass
class DiffiTModelBundle:
    checkpoint_name: str
    checkpoint_path: str
    image_size: int
    latent_size: int
    default_cfg_scale: float
    model: torch.nn.Module
    diffusion_cache: dict[int, Any] = field(default_factory=dict)
    current_device: str = "cpu"
    precision: str = "fp32"


@dataclass
class DiffiTConditioning:
    class_id: int
    label: str


@dataclass
class DiffiTLatent:
    tensor: torch.Tensor
    image_size: int
    batch_size: int
    seed: int


_MODEL_CACHE: dict[tuple[str, int], DiffiTModelBundle] = {}
_VAE_CACHE: dict[tuple[str, str], Any] = {}


def _lazy_import_diffit_model():
    try:
        module = importlib.import_module(".vendor.diffit_upstream.diffit", package=__package__)
    except ImportError as exc:
        raise RuntimeError(
            "DiffiT model import failed. Install this node pack's requirements first "
            "(missing dependency is usually timm)."
        ) from exc
    return module


def _lazy_import_autoencoder():
    try:
        module = importlib.import_module("diffusers.models")
    except ImportError as exc:
        raise RuntimeError(
            "DiffiT VAE decode requires diffusers. Install this node pack's requirements first."
        ) from exc
    return module.AutoencoderKL


def _move_module(module: torch.nn.Module, device: torch.device) -> torch.nn.Module:
    module.to(device)
    module.eval()
    return module


def _get_module_dtype(module: torch.nn.Module) -> torch.dtype:
    try:
        return next(module.parameters()).dtype
    except StopIteration:
        return torch.float32


def _make_progress_bar(total_steps: int, node_id=None):
    try:
        import comfy.utils  # type: ignore

        return comfy.utils.ProgressBar(total_steps, node_id=str(node_id) if node_id is not None else None)
    except Exception:
        return None


def _ensure_bundle_device(bundle: DiffiTModelBundle, device: torch.device) -> DiffiTModelBundle:
    if bundle.current_device != str(device):
        LOGGER.info("Moving DiffiT model to %s", device)
        _move_module(bundle.model, device)
        bundle.current_device = str(device)
    return bundle


def _ensure_bundle_precision(bundle: DiffiTModelBundle, use_fp16: bool) -> DiffiTModelBundle:
    target_precision = "fp16" if use_fp16 else "fp32"
    target_dtype = torch.float16 if use_fp16 else torch.float32
    current_dtype = _get_module_dtype(bundle.model)
    if current_dtype != target_dtype:
        LOGGER.info("Switching DiffiT model precision to %s", target_precision)
        if use_fp16 and hasattr(bundle.model, "convert_to_fp16"):
            bundle.model.convert_to_fp16()
        elif not use_fp16 and hasattr(bundle.model, "convert_to_fp32"):
            bundle.model.convert_to_fp32()
        else:
            bundle.model.to(dtype=target_dtype)
    bundle.precision = target_precision
    bundle.model.eval()
    return bundle


def _build_diffusion(steps: int):
    config = diffusion_defaults()
    config["timestep_respacing"] = str(int(steps))
    return create_diffusion(**config)


def _get_cached_diffusion(bundle: DiffiTModelBundle, steps: int):
    if steps not in bundle.diffusion_cache:
        bundle.diffusion_cache[steps] = _build_diffusion(steps)
    return bundle.diffusion_cache[steps]


def _resolve_cfg_eps_channels(model_bundle: DiffiTModelBundle, cfg_channel_mode: str) -> int:
    if cfg_channel_mode == "latent_channels":
        return int(getattr(model_bundle.model, "in_channels", 4))
    if cfg_channel_mode == "upstream_repo_3ch":
        return 3
    raise ValueError(
        f"Unsupported cfg_channel_mode '{cfg_channel_mode}'. "
        "Use 'latent_channels' or 'upstream_repo_3ch'."
    )


def _download_checkpoint_from_hf(checkpoint_name: str) -> Path:
    if checkpoint_name not in {"diffit_256.safetensors", "diffit_512.safetensors"}:
        raise FileNotFoundError(
            f"Checkpoint {checkpoint_name} was not found locally and is not one of the official DiffiT filenames."
        )

    try:
        from huggingface_hub import hf_hub_download
        from huggingface_hub.utils import HfHubHTTPError
    except ImportError as exc:
        raise RuntimeError(
            "Downloading DiffiT weights requires huggingface_hub. Install this node pack's requirements first."
        ) from exc

    target_dir = get_preferred_diffit_download_dir()
    LOGGER.info("Downloading %s into %s", checkpoint_name, target_dir)
    try:
        local_path = hf_hub_download(
            repo_id="nvidia/DiffiT",
            filename=checkpoint_name,
            local_dir=str(target_dir),
        )
    except HfHubHTTPError as exc:
        raise RuntimeError(
            "Failed to download the official DiffiT checkpoint from Hugging Face. "
            "Check your network connection and any Hugging Face auth configuration."
        ) from exc
    except Exception as exc:
        raise RuntimeError(f"Unexpected DiffiT download failure for {checkpoint_name}.") from exc

    return Path(local_path)


def load_diffit_model(
    checkpoint_name: str,
    checkpoint_path: str = "",
    image_size_override: str = "auto",
    download_if_missing: bool = True,
    device: str = "auto",
) -> DiffiTModelBundle:
    resolved_path = resolve_diffit_checkpoint_path(checkpoint_name, checkpoint_path)
    if resolved_path is None and download_if_missing:
        resolved_path = _download_checkpoint_from_hf(checkpoint_name)

    if resolved_path is None:
        target_dir = get_preferred_diffit_download_dir()
        raise FileNotFoundError(
            f"DiffiT checkpoint {checkpoint_name} was not found. Place it in {target_dir} "
            f"or in {target_dir.parent.parent / 'Packages' / 'ComfyUI-rect' / 'models' / 'diffusion_models'}."
        )

    state_dict = load_state_dict_file(resolved_path)
    inferred_size = infer_image_size_from_state_dict(state_dict)
    image_size = inferred_size if image_size_override == "auto" else validate_image_size(int(image_size_override))
    cache_key = (str(resolved_path.resolve()), image_size)

    bundle = _MODEL_CACHE.get(cache_key)
    if bundle is None:
        diffit_module = _lazy_import_diffit_model()
        latent_size = latent_size_for_image(image_size)
        model = diffit_module.Diffit(input_size=latent_size)
        load_result = model.load_state_dict(state_dict, strict=True)
        if getattr(load_result, "missing_keys", None) or getattr(load_result, "unexpected_keys", None):
            LOGGER.warning(
                "DiffiT checkpoint load reported mismatched keys. Missing=%s Unexpected=%s",
                getattr(load_result, "missing_keys", []),
                getattr(load_result, "unexpected_keys", []),
            )
        model.eval()
        bundle = DiffiTModelBundle(
            checkpoint_name=resolved_path.name,
            checkpoint_path=str(resolved_path),
            image_size=image_size,
            latent_size=latent_size,
            default_cfg_scale=default_cfg_scale_for_image_size(image_size),
            model=model,
            diffusion_cache={250: _build_diffusion(250)},
            current_device="cpu",
        )
        _MODEL_CACHE[cache_key] = bundle

    target_device = pick_device(device)
    return _ensure_bundle_device(bundle, target_device)


def create_class_conditioning(
    class_id: int,
    label_choice: str = "<use class_id>",
) -> DiffiTConditioning:
    resolved_class_id = int(class_id)
    resolved_label = f"class_{resolved_class_id}"

    if label_choice and label_choice != "<use class_id>":
        prefix = label_choice.split(":", 1)[0].strip()
        try:
            resolved_class_id = int(prefix)
        except ValueError:
            raise ValueError(f"Unable to parse ImageNet class id from label choice: {label_choice}")
        resolved_label = label_choice.split(":", 1)[1].strip() if ":" in label_choice else label_choice

    if not 0 <= resolved_class_id < NUM_CLASSES:
        raise ValueError(f"DiffiT class id must be between 0 and 999. Received {resolved_class_id}.")

    return DiffiTConditioning(class_id=resolved_class_id, label=resolved_label)


def create_empty_latent(batch_size: int, image_size: int, seed: int) -> DiffiTLatent:
    image_size = validate_image_size(int(image_size))
    batch_size = int(batch_size)
    if batch_size < 1:
        raise ValueError("DiffiT batch_size must be at least 1.")

    latent_size = latent_size_for_image(image_size)
    generator = torch.Generator(device="cpu")
    generator.manual_seed(int(seed) & 0xFFFFFFFFFFFFFFFF)
    tensor = torch.randn((batch_size, 4, latent_size, latent_size), generator=generator, dtype=torch.float32)
    return DiffiTLatent(tensor=tensor, image_size=image_size, batch_size=batch_size, seed=int(seed))


def sample_diffit(
    model_bundle: DiffiTModelBundle,
    conditioning: DiffiTConditioning,
    latent: DiffiTLatent,
    steps: int = 250,
    cfg_scale: float = 0.0,
    seed: int = 0,
    use_ddim: bool = False,
    use_fp16: bool = True,
    scale_pow: float = 4.0,
    cfg_enabled: bool = True,
    cfg_channel_mode: str = "latent_channels",
    node_id=None,
) -> DiffiTLatent:
    if int(steps) < 2:
        raise ValueError("DiffiT sampling requires at least 2 steps.")

    if latent.image_size != model_bundle.image_size:
        raise ValueError(
            f"Latent image_size={latent.image_size} does not match loaded DiffiT model image_size={model_bundle.image_size}."
        )

    device = torch.device(model_bundle.current_device)
    bundle = _ensure_bundle_device(model_bundle, device)
    fp16_requested = bool(use_fp16 and device.type == "cuda")
    _ensure_bundle_precision(bundle, fp16_requested)
    diffusion = _get_cached_diffusion(bundle, int(steps))
    latent_tensor = latent.tensor.to(device=device, dtype=torch.float32, non_blocking=device.type == "cuda")

    model_kwargs: dict[str, Any] = {}
    sample_model = bundle.model.forward
    if cfg_enabled:
        z = torch.cat([latent_tensor, latent_tensor], dim=0)
        classes = torch.full((latent.batch_size,), conditioning.class_id, device=device, dtype=torch.long)
        classes_null = torch.full((latent.batch_size,), NUM_CLASSES, device=device, dtype=torch.long)
        model_kwargs["y"] = torch.cat([classes, classes_null], dim=0)
        model_kwargs["cfg_scale"] = float(cfg_scale) if float(cfg_scale) > 0 else bundle.default_cfg_scale
        model_kwargs["diffusion_steps"] = diffusion_defaults()["diffusion_steps"]
        model_kwargs["scale_pow"] = float(scale_pow)
        model_kwargs["cfg_eps_channels"] = _resolve_cfg_eps_channels(bundle, cfg_channel_mode)
        sample_model = bundle.model.forward_with_cfg
    else:
        z = latent_tensor
        model_kwargs["y"] = torch.full(
            (latent.batch_size,),
            conditioning.class_id,
            device=device,
            dtype=torch.long,
        )

    sample_fn = diffusion.ddim_sample_loop if use_ddim else diffusion.p_sample_loop
    progress_bar = _make_progress_bar(int(steps), node_id=node_id)

    if device.type == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True

    def run_sampler(autocast_enabled: bool) -> torch.Tensor:
        generator = make_generator(device, int(seed))
        step_callback = None
        if progress_bar is not None:
            step_callback = lambda current, total: progress_bar.update_absolute(current, total=total)
        with torch.inference_mode():
            with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=autocast_enabled):
                return sample_fn(
                    sample_model,
                    z.shape,
                    noise=z,
                    clip_denoised=False,
                    progress=False,
                    model_kwargs=model_kwargs,
                    device=device,
                    generator=generator,
                    step_callback=step_callback,
                )

    def run_sampler_with_fallback() -> torch.Tensor:
        autocast_enabled = fp16_requested
        try:
            sampled_tensor = run_sampler(autocast_enabled)
        except RuntimeError as exc:
            if not fp16_requested:
                raise
            LOGGER.warning("DiffiT fp16 sampler failed (%s). Retrying in fp32.", exc)
            _ensure_bundle_precision(bundle, False)
            if progress_bar is not None:
                progress_bar.update_absolute(0, total=int(steps))
            return run_sampler(False)

        if fp16_requested and not torch.isfinite(sampled_tensor).all():
            LOGGER.warning("DiffiT fp16 sampler produced non-finite values. Retrying in fp32.")
            _ensure_bundle_precision(bundle, False)
            if progress_bar is not None:
                progress_bar.update_absolute(0, total=int(steps))
            return run_sampler(False)
        return sampled_tensor

    sampled = run_sampler_with_fallback()

    if cfg_enabled:
        sampled, _ = sampled.chunk(2, dim=0)

    if not torch.isfinite(sampled).all():
        raise RuntimeError("DiffiT sampler produced non-finite latent values.")

    return DiffiTLatent(
        tensor=sampled.detach().cpu().float(),
        image_size=bundle.image_size,
        batch_size=latent.batch_size,
        seed=int(seed),
    )


def sample_diffit_official(
    model_bundle: DiffiTModelBundle,
    conditioning: DiffiTConditioning,
    latent: DiffiTLatent,
    steps: int = 250,
    cfg_scale: float = 0.0,
    use_ddim: bool = False,
    use_fp16: bool = False,
    scale_pow: float = 4.0,
    cfg_cond: bool = True,
    clip_denoised: bool = False,
    node_id=None,
) -> DiffiTLatent:
    """
    Follow the released NVlabs sample.py path as closely as practical in ComfyUI.

    This intentionally preserves the public repo's CFG behavior by applying
    guidance through ``forward_with_cfg`` with 3 guided channels.
    """
    if int(steps) < 2:
        raise ValueError("DiffiT official sampling requires at least 2 steps.")

    if latent.image_size != model_bundle.image_size:
        raise ValueError(
            f"Latent image_size={latent.image_size} does not match loaded DiffiT model image_size={model_bundle.image_size}."
        )

    device = torch.device(model_bundle.current_device)
    if use_fp16 and device.type != "cuda":
        raise RuntimeError("DiffiT official fp16 mode requires a CUDA device.")

    bundle = _ensure_bundle_device(model_bundle, device)
    _ensure_bundle_precision(bundle, bool(use_fp16 and device.type == "cuda"))
    diffusion = _get_cached_diffusion(bundle, int(steps))

    model_dtype = _get_module_dtype(bundle.model)
    latent_tensor = latent.tensor.to(device=device, dtype=model_dtype, non_blocking=device.type == "cuda")

    model_kwargs: dict[str, Any] = {}
    if cfg_cond:
        z = torch.cat([latent_tensor, latent_tensor], dim=0)
        classes = torch.full((latent.batch_size,), conditioning.class_id, device=device, dtype=torch.long)
        classes_null = torch.full((latent.batch_size,), NUM_CLASSES, device=device, dtype=torch.long)
        model_kwargs["y"] = torch.cat([classes, classes_null], dim=0)
        model_kwargs["cfg_scale"] = float(cfg_scale) if float(cfg_scale) > 0 else bundle.default_cfg_scale
        model_kwargs["diffusion_steps"] = diffusion_defaults()["diffusion_steps"]
        model_kwargs["scale_pow"] = float(scale_pow)
    else:
        z = latent_tensor
        model_kwargs["y"] = torch.full(
            (latent.batch_size,),
            conditioning.class_id,
            device=device,
            dtype=torch.long,
        )

    sample_fn = diffusion.ddim_sample_loop if use_ddim else diffusion.p_sample_loop
    progress_bar = _make_progress_bar(int(steps), node_id=node_id)

    if device.type == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True

    step_callback = None
    if progress_bar is not None:
        step_callback = lambda current, total: progress_bar.update_absolute(current, total=total)

    with torch.inference_mode():
        sampled = sample_fn(
            bundle.model.forward_with_cfg,
            z.shape,
            z,
            clip_denoised=bool(clip_denoised),
            progress=False,
            model_kwargs=model_kwargs,
            device=device,
            step_callback=step_callback,
        )

    if cfg_cond:
        sampled, _ = sampled.chunk(2, dim=0)

    if not torch.isfinite(sampled).all():
        raise RuntimeError(
            "DiffiT official sampler produced non-finite latent values. "
            "This follows the released repo path closely; disable fp16 or use the non-official sampler if needed."
        )

    return DiffiTLatent(
        tensor=sampled.detach().cpu().float(),
        image_size=bundle.image_size,
        batch_size=latent.batch_size,
        seed=latent.seed,
    )


def _load_vae(variant: str, device: torch.device):
    variant = variant.lower()
    if variant not in {"ema", "mse"}:
        raise ValueError(f"Unsupported DiffiT VAE variant '{variant}'. Use 'ema' or 'mse'.")

    cache_key = (variant, str(device))
    cached = _VAE_CACHE.get(cache_key)
    if cached is not None:
        return cached

    AutoencoderKL = _lazy_import_autoencoder()
    repo_id = f"stabilityai/sd-vae-ft-{variant}"
    cache_dir = get_hf_cache_dir(f"sd-vae-ft-{variant}")
    try:
        vae = AutoencoderKL.from_pretrained(repo_id, cache_dir=str(cache_dir))
    except Exception as exc:
        raise RuntimeError(
            "Failed to load the DiffiT VAE from Hugging Face. "
            "Check your network connection, local Hugging Face cache, and any auth configuration."
        ) from exc

    vae = _move_module(vae, device)
    if hasattr(vae, "enable_slicing"):
        vae.enable_slicing()
    _VAE_CACHE[cache_key] = vae
    return vae


def decode_diffit_latent(
    latent: DiffiTLatent,
    vae_variant: str = "ema",
    use_fp16: bool = True,
    device: str = "auto",
) -> torch.Tensor:
    target_device = pick_device(device)
    vae = _load_vae(vae_variant, target_device)
    decode_fp16 = bool(use_fp16 and target_device.type == "cuda")
    vae.to(dtype=torch.float16 if decode_fp16 else torch.float32)
    latent_tensor = latent.tensor.to(target_device, non_blocking=target_device.type == "cuda")

    def run_decode(autocast_enabled: bool) -> torch.Tensor:
        with torch.inference_mode():
            with torch.autocast(
                device_type=target_device.type,
                dtype=torch.float16,
                enabled=autocast_enabled,
            ):
                return vae.decode(latent_tensor / 0.18215).sample

    def run_decode_with_fallback() -> torch.Tensor:
        try:
            decoded_tensor = run_decode(decode_fp16)
        except RuntimeError as exc:
            if not decode_fp16:
                raise
            LOGGER.warning("DiffiT fp16 VAE decode failed (%s). Retrying in fp32.", exc)
            vae.to(dtype=torch.float32)
            return run_decode(False)

        if decode_fp16 and not torch.isfinite(decoded_tensor).all():
            LOGGER.warning("DiffiT fp16 VAE decode produced non-finite values. Retrying in fp32.")
            vae.to(dtype=torch.float32)
            return run_decode(False)
        return decoded_tensor

    decoded = run_decode_with_fallback()

    if not torch.isfinite(decoded).all():
        raise RuntimeError("DiffiT VAE decode produced non-finite image values.")

    return to_comfy_image(decoded)


def decode_diffit_latent_official(
    latent: DiffiTLatent,
    vae_variant: str = "ema",
    device: str = "auto",
) -> torch.Tensor:
    """
    Follow the released sample.py VAE path: decode in the default VAE dtype.
    """
    target_device = pick_device(device)
    vae = _load_vae(vae_variant, target_device)
    vae.to(dtype=torch.float32)
    latent_tensor = latent.tensor.to(
        target_device,
        dtype=torch.float32,
        non_blocking=target_device.type == "cuda",
    )

    with torch.inference_mode():
        decoded = vae.decode(latent_tensor / 0.18215).sample

    if not torch.isfinite(decoded).all():
        raise RuntimeError("DiffiT official VAE decode produced non-finite image values.")

    return to_comfy_image(decoded)


def generate_diffit_image(
    model_bundle: DiffiTModelBundle,
    conditioning: DiffiTConditioning,
    batch_size: int,
    latent_seed: int,
    steps: int = 250,
    cfg_scale: float = 0.0,
    sampler_seed: int = 0,
    use_ddim: bool = False,
    use_fp16: bool = True,
    scale_pow: float = 4.0,
    cfg_enabled: bool = True,
    cfg_channel_mode: str = "latent_channels",
    vae_variant: str = "ema",
    node_id=None,
) -> tuple[torch.Tensor, DiffiTLatent]:
    latent = create_empty_latent(batch_size=batch_size, image_size=model_bundle.image_size, seed=latent_seed)
    sampled = sample_diffit(
        model_bundle=model_bundle,
        conditioning=conditioning,
        latent=latent,
        steps=steps,
        cfg_scale=cfg_scale,
        seed=sampler_seed,
        use_ddim=use_ddim,
        use_fp16=use_fp16,
        scale_pow=scale_pow,
        cfg_enabled=cfg_enabled,
        cfg_channel_mode=cfg_channel_mode,
        node_id=node_id,
    )
    image = decode_diffit_latent(sampled, vae_variant=vae_variant, use_fp16=use_fp16, device=model_bundle.current_device)
    return image, sampled


def generate_diffit_image_official(
    model_bundle: DiffiTModelBundle,
    conditioning: DiffiTConditioning,
    batch_size: int,
    latent_seed: int,
    steps: int = 250,
    cfg_scale: float = 0.0,
    use_ddim: bool = False,
    use_fp16: bool = False,
    scale_pow: float = 4.0,
    cfg_cond: bool = True,
    clip_denoised: bool = False,
    vae_variant: str = "ema",
    node_id=None,
) -> tuple[torch.Tensor, DiffiTLatent]:
    latent = create_empty_latent(batch_size=batch_size, image_size=model_bundle.image_size, seed=latent_seed)
    sampled = sample_diffit_official(
        model_bundle=model_bundle,
        conditioning=conditioning,
        latent=latent,
        steps=steps,
        cfg_scale=cfg_scale,
        use_ddim=use_ddim,
        use_fp16=use_fp16,
        scale_pow=scale_pow,
        cfg_cond=cfg_cond,
        clip_denoised=clip_denoised,
        node_id=node_id,
    )
    image = decode_diffit_latent_official(
        sampled,
        vae_variant=vae_variant,
        device=model_bundle.current_device,
    )
    return image, sampled