"""ComfyUI node definitions for the official single-node DiffiT pipeline."""

from __future__ import annotations

from .diffit_runtime import (
    create_class_conditioning,
    generate_diffit_image_official,
    load_diffit_model,
)
from .utils import LOGGER, list_available_diffit_checkpoints


CATEGORY = "NVIDIA/DiffiT"


class DiffiTOfficialImageNode:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "checkpoint_name": (list_available_diffit_checkpoints(),),
                "download_if_missing": ("BOOLEAN", {"default": True}),
                "device": (["auto", "cuda", "cpu"],),
                "class_id": ("INT", {"default": 207, "min": 0, "max": 999}),
                "batch_size": ("INT", {"default": 1, "min": 1, "max": 64}),
                "seed": ("INT", {"default": 12345, "min": 0, "max": 0x7FFFFFFFFFFFFFFF}),
                "steps": ("INT", {"default": 250, "min": 2, "max": 1000}),
                "cfg_scale": ("FLOAT", {"default": 4.4, "min": 0.0, "max": 32.0, "step": 0.01}),
                "cfg_cond": ("BOOLEAN", {"default": True}),
                "use_ddim": ("BOOLEAN", {"default": False}),
                "scale_pow": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 16.0, "step": 0.1}),
                "clip_denoised": ("BOOLEAN", {"default": False}),
                "vae_variant": (["ema", "mse"],),
            },
            "hidden": {"unique_id": "UNIQUE_ID"},
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "generate"
    CATEGORY = CATEGORY

    def generate(
        self,
        checkpoint_name,
        download_if_missing,
        device,
        class_id,
        batch_size,
        seed,
        steps,
        cfg_scale,
        cfg_cond,
        use_ddim,
        scale_pow,
        clip_denoised,
        vae_variant,
        unique_id=None,
    ):
        try:
            model = load_diffit_model(
                checkpoint_name=checkpoint_name,
                checkpoint_path="",
                image_size_override="auto",
                download_if_missing=download_if_missing,
                device=device,
            )
            conditioning = create_class_conditioning(class_id, "<use class_id>")
            image, _sampled = generate_diffit_image_official(
                model_bundle=model,
                conditioning=conditioning,
                batch_size=batch_size,
                latent_seed=seed,
                steps=steps,
                cfg_scale=cfg_scale,
                cfg_cond=cfg_cond,
                use_ddim=use_ddim,
                use_fp16=False,
                scale_pow=scale_pow,
                clip_denoised=clip_denoised,
                vae_variant=vae_variant,
                node_id=unique_id,
            )
            return (image,)
        except Exception as exc:
            LOGGER.exception("DiffiT official image generation failed.")
            raise RuntimeError(f"DiffiT Official Image failed: {exc}") from exc


NODE_CLASS_MAPPINGS = {
    "DiffiTOfficialImage": DiffiTOfficialImageNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "DiffiTOfficialImage": "DiffiT Official Image",
}