import atexit

import torch
import nvvfx
import comfy.model_management
from enum import Enum
from typing import TypedDict
from typing_extensions import override

from comfy_api.latest import ComfyExtension, io


# nvidia-vfx 0.1.0.1 can block while destroying a VideoSuperRes context on
# Windows. Keep loaded effects alive and reuse them for matching jobs. This also
# avoids reloading the RTX model for every video chunk in a ComfyUI workflow.
_EFFECT_CACHE: dict[tuple[int, str, int, int], nvvfx.VideoSuperRes] = {}


def _skip_blocking_effect_teardown() -> None:
    # The Windows process will reclaim the native CUDA allocations. Marking the
    # objects closed prevents Effect.__del__ from entering a known-blocking
    # destroy() call during interpreter shutdown.
    for effect in _EFFECT_CACHE.values():
        effect._closed = True


atexit.register(_skip_blocking_effect_teardown)


class UpscaleType(str, Enum):
    SCALE_BY = "scale by multiplier"
    TARGET_DIMENSIONS = "target dimensions"


class RTXVideoSuperResolution(io.ComfyNode):
    class UpscaleTypedDict(TypedDict):
        resize_type: UpscaleType
        scale: float
        width: int
        height: int

    @classmethod
    def define_schema(cls):
        return io.Schema(
            node_id="RTXVideoSuperResolution",
            display_name="RTX Video Super Resolution",
            category="image/upscaling",
            search_aliases=["rtx", "nvidia", "upscale", "super resolution", "vsr"],
            inputs=[
                io.Image.Input("images"),
                io.DynamicCombo.Input(
                    "resize_type",
                    tooltip="Choose to scale by a multiplier or to exact target dimensions.",
                    options=[
                        io.DynamicCombo.Option(UpscaleType.SCALE_BY, [
                            io.Float.Input("scale", default=2.0, min=1.0, max=4.0, step=0.01, tooltip="Scale factor (e.g., 2.0 doubles the size)."),
                        ]),
                        io.DynamicCombo.Option(UpscaleType.TARGET_DIMENSIONS, [
                            io.Int.Input("width", default=1920, min=64, max=8192, step=8, tooltip="Target width in pixels."),
                            io.Int.Input("height", default=1080, min=64, max=8192, step=8, tooltip="Target height in pixels.")
                        ])
                    ],
                ),
                io.Combo.Input(
                    "quality",
                    options=[
                        "LOW",
                        "MEDIUM",
                        "HIGH",
                        "ULTRA",
                        "HIGHBITRATE_HIGH",
                        "HIGHBITRATE_ULTRA",
                    ],
                    default="HIGHBITRATE_HIGH",
                ),
            ],
            outputs=[
                io.Image.Output("upscaled_images"),
            ],
        )

    @classmethod
    def execute(cls, images: torch.Tensor, resize_type: UpscaleTypedDict, quality: str) -> io.NodeOutput:
        b, h, w, c = images.shape

        selected_type = resize_type["resize_type"]
        if selected_type == UpscaleType.SCALE_BY:
            scale = resize_type["scale"]
            output_width = int(w * scale)
            output_height = int(h * scale)
        elif selected_type == UpscaleType.TARGET_DIMENSIONS:
            output_width = resize_type["width"]
            output_height = resize_type["height"]
        else:
            raise ValueError(f"Unsupported resize type: {selected_type}")

        output_width = max(8, round(output_width / 8) * 8)
        output_height = max(8, round(output_height / 8) * 8)

        MAX_PIXELS = 1024 * 1024 * 16

        out_pixels = output_width * output_height
        batch_size = max(1, MAX_PIXELS // out_pixels)

        quality_mapping = {
            "LOW": nvvfx.effects.QualityLevel.LOW,
            "MEDIUM": nvvfx.effects.QualityLevel.MEDIUM,
            "HIGH": nvvfx.effects.QualityLevel.HIGH,
            "ULTRA": nvvfx.effects.QualityLevel.ULTRA,
            "HIGHBITRATE_HIGH": nvvfx.effects.QualityLevel.HIGHBITRATE_HIGH,
            "HIGHBITRATE_ULTRA": nvvfx.effects.QualityLevel.HIGHBITRATE_ULTRA,
        }
        selected_quality = quality_mapping.get(
            quality, nvvfx.effects.QualityLevel.HIGHBITRATE_HIGH
        )

        device_index = images.device.index
        if device_index is None:
            device_index = torch.cuda.current_device()
        effect_key = (device_index, quality, output_width, output_height)
        sr = _EFFECT_CACHE.get(effect_key)
        if sr is None:
            # RTX VSR is normally the final stage of a large video workflow.
            # Explicitly offload H3/VAEs before nvvfx reserves its own CUDA
            # context; otherwise a 24 GB card can fail despite the decoded
            # frames themselves living on the CPU intermediate device.
            comfy.model_management.unload_all_models()
            comfy.model_management.soft_empty_cache()
            sr = nvvfx.VideoSuperRes(selected_quality, device=device_index)
            sr.output_width = output_width
            sr.output_height = output_height
            sr.load()
            _EFFECT_CACHE[effect_key] = sr

        out_tensor = torch.empty((images.shape[0], output_height, output_width, c), device=images.device, dtype=images.dtype)
        for i in range(0, images.shape[0], batch_size):
            batch = images[i:i + batch_size]

            batch_cuda = batch.to(device=f"cuda:{device_index}").permute(0, 3, 1, 2).float().contiguous()

            for j in range(batch_cuda.shape[0]):
                input_frame = batch_cuda[j]
                dlpack_out = sr.run(input_frame).image
                out_tensor[i + j: i + j + 1] = torch.from_dlpack(dlpack_out).movedim(0, -1).unsqueeze(0)

        return io.NodeOutput(out_tensor)


class NVVFXVideoExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[io.ComfyNode]]:
        return [
            RTXVideoSuperResolution,
        ]


async def comfy_entrypoint() -> NVVFXVideoExtension:
    return NVVFXVideoExtension()

# hack so registry picks up the node name
if False:
    NODE_CLASS_MAPPINGS = {"RTXVideoSuperResolution": RTXVideoSuperResolution}
