"""The two ComfyUI nodes exposed by Better NormalCrafter Wrapper."""

import comfy.model_management
import comfy.utils

from .normalcrafter_clean.inference import estimate_progress_steps, run_inference
from .normalcrafter_clean.model import (
    base_model_sources,
    load_models,
    normalcrafter_sources,
)


class NormalCrafterLoader:
    """Choose and load the NormalCrafter and SVD model components."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "normalcrafter_model": (
                    normalcrafter_sources(),
                    {
                        "default": "HF: Yanrui95/NormalCrafter",
                        "tooltip": (
                            "NormalCrafter model containing the UNet and VAE. "
                            "Local models are discovered from every ComfyUI "
                            "'normalcrafter' model path, including paths from "
                            "extra_model_paths.yaml. Selecting an HF entry downloads "
                            "the required files into a registered normalcrafter path."
                        ),
                    },
                ),
                "base_model": (
                    base_model_sources(),
                    {
                        "default": "HF: stabilityai/stable-video-diffusion-img2vid-xt",
                        "tooltip": (
                            "Stable Video Diffusion base components used by "
                            "NormalCrafter: CLIP image encoder, feature extractor "
                            "and scheduler. Local models are discovered from the "
                            "same registered ComfyUI 'normalcrafter' paths."
                        ),
                    },
                ),
            }
        }

    RETURN_TYPES = ("NORMALCRAFTER_MODEL",)
    RETURN_NAMES = ("model",)
    OUTPUT_TOOLTIPS = (
        "Loaded NormalCrafter model bundle. Heavy components stay on CPU until inference needs them.",
    )
    FUNCTION = "load"
    CATEGORY = "NormalCrafter"
    DESCRIPTION = (
        "Loads the NormalCrafter UNet/VAE and the required Stable Video Diffusion "
        "base components. Model discovery follows ComfyUI's registered "
        "'normalcrafter' paths, including extra_model_paths.yaml."
    )

    def load(self, normalcrafter_model: str, base_model: str):
        device = comfy.model_management.get_torch_device()
        return (
            load_models(
                device=device,
                normalcrafter_source=normalcrafter_model,
                base_model_source=base_model,
            ),
        )


class NormalCrafterGenerate:
    """Run NormalCrafter with explicit temporal and VRAM controls."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": (
                    "NORMALCRAFTER_MODEL",
                    {
                        "tooltip": (
                            "NormalCrafter model bundle from the NormalCrafter - Load node."
                        )
                    },
                ),
                "images": (
                    "IMAGE",
                    {
                        "tooltip": (
                            "Input image batch or video frames in temporal order. "
                            "The output contains one normal map per input frame."
                        )
                    },
                ),
                "max_resolution": (
                    "INT",
                    {
                        "default": 1024,
                        "min": 256,
                        "max": 2048,
                        "step": 64,
                        "tooltip": (
                            "Maximum length of the input's longest side during inference. "
                            "Images are never enlarged. Lower values reduce VRAM use and "
                            "increase speed, at the cost of spatial detail. Output is "
                            "resized back to the original input resolution."
                        ),
                    },
                ),
                "window_size": (
                    "INT",
                    {
                        "default": 14,
                        "min": 2,
                        "max": 64,
                        "step": 1,
                        "tooltip": (
                            "Number of consecutive frames processed by the temporal UNet "
                            "at once. Larger windows provide more temporal context but use "
                            "more VRAM. NormalCrafter's released default is 14."
                        ),
                    },
                ),
                "step_size": (
                    "INT",
                    {
                        "default": 10,
                        "min": 1,
                        "max": 64,
                        "step": 1,
                        "tooltip": (
                            "How many frames the temporal window advances each pass. Must "
                            "be <= window_size. Smaller values create more overlap and "
                            "usually smoother temporal blending, but require more UNet "
                            "passes. With window_size 14, the released default is 10 "
                            "(4-frame regular overlap)."
                        ),
                    },
                ),
                "chunk_size": (
                    "INT",
                    {
                        "default": 7,
                        "min": 1,
                        "max": 64,
                        "step": 1,
                        "tooltip": (
                            "Batch size used for CLIP encoding, VAE encoding and VAE "
                            "decoding. Lower values reduce peak VRAM; higher values can "
                            "improve throughput if enough VRAM is available. This does "
                            "not change the temporal UNet window size."
                        ),
                    },
                ),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("normal_maps",)
    OUTPUT_TOOLTIPS = (
        "Per-frame RGB normal maps in ComfyUI IMAGE format, restored to the input resolution.",
    )
    FUNCTION = "generate"
    CATEGORY = "NormalCrafter"
    DESCRIPTION = (
        "Runs NormalCrafter in four stages: CLIP encode, VAE encode, temporal UNet "
        "inference and VAE decode. Full-video intermediates stay on CPU and only "
        "the active component/chunk/window is moved to the inference device."
    )

    def generate(
        self,
        model,
        images,
        max_resolution: int,
        window_size: int,
        step_size: int,
        chunk_size: int,
    ):
        device = comfy.model_management.get_torch_device()
        total_steps = estimate_progress_steps(
            int(images.shape[0]),
            window_size=window_size,
            step_size=step_size,
            chunk_size=chunk_size,
        )
        progress_bar = comfy.utils.ProgressBar(total_steps)

        def progress() -> None:
            progress_bar.update(1)

        output = run_inference(
            model,
            images,
            device=device,
            max_resolution=max_resolution,
            window_size=window_size,
            step_size=step_size,
            chunk_size=chunk_size,
            progress=progress,
        )

        # The model lifecycle already returned all modules to CPU. This only lets
        # ComfyUI release allocator blocks that are now unused.
        comfy.model_management.soft_empty_cache()
        return (output,)


NODE_CLASS_MAPPINGS = {
    "NormalCrafterLoader": NormalCrafterLoader,
    "NormalCrafterGenerate": NormalCrafterGenerate,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "NormalCrafterLoader": "NormalCrafter - Load",
    "NormalCrafterGenerate": "NormalCrafter - Generate Normals",
}