from __future__ import annotations

import torch

if __package__:
    from .mask_ops import GRID_MODES, H3_PIXEL_CELL, quantize_h3_pixel_mask
    from .media_ops import H3_VIDEO_FPS, floor_h3_frame_count, trim_audio_to_frames
    from .runtime import install_model_wrappers
else:
    from mask_ops import GRID_MODES, H3_PIXEL_CELL, quantize_h3_pixel_mask
    from media_ops import H3_VIDEO_FPS, floor_h3_frame_count, trim_audio_to_frames
    from runtime import install_model_wrappers


class MiniMaxH3TrimSourceAV:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "frames": (
                    "IMAGE",
                    {
                        "tooltip": (
                            "Source video frame batch. Trailing frames are dropped to the largest valid "
                            "MiniMax H3 17k+5 frame count. Frames are never padded or resized."
                        )
                    },
                ),
            },
            "optional": {
                "audio": (
                    "AUDIO",
                    {
                        "tooltip": (
                            "Optional synchronized source audio. It is trimmed to the duration of the "
                            "returned frame batch at H3's fixed 24 fps."
                        )
                    },
                )
            },
        }

    RETURN_TYPES = ("IMAGE", "AUDIO", "INT", "STRING")
    RETURN_NAMES = ("trimmed_frames", "trimmed_audio", "h3_length", "trim_info")
    FUNCTION = "trim"
    CATEGORY = "image/video/minimax"
    DESCRIPTION = (
        "Drops trailing source frames to the largest valid MiniMax H3 17k+5 length and, "
        "when connected, trims synchronized audio to the same 24 fps duration."
    )

    def trim(self, frames, audio=None):
        if frames.ndim != 4 or frames.shape[-1] < 3:
            raise ValueError(
                f"H3 AV Trim expects ComfyUI IMAGE frames [T,H,W,C], got {list(frames.shape)}."
            )

        source_count = int(frames.shape[0])
        h3_length = floor_h3_frame_count(source_count)
        trimmed_frames = frames[:h3_length]
        trimmed_audio, target_samples = trim_audio_to_frames(audio, h3_length)

        dropped = source_count - h3_length
        duration = h3_length / H3_VIDEO_FPS
        audio_text = "no audio connected"
        if trimmed_audio is not None:
            actual_samples = int(trimmed_audio["waveform"].shape[-1])
            sample_rate = int(trimmed_audio["sample_rate"])
            audio_text = f"audio {actual_samples}/{target_samples} samples at {sample_rate} Hz"
            if actual_samples < target_samples:
                audio_text += " (source audio was shorter; no silence was added)"
        info = (
            f"H3 source trim: {source_count} -> {h3_length} frames; dropped {dropped}; "
            f"duration {duration:.3f}s at {H3_VIDEO_FPS} fps; {audio_text}."
        )
        return (trimmed_frames, trimmed_audio, h3_length, info)


class MiniMaxH3PerRowMaskPatch:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": (
                    "MODEL",
                    {
                        "tooltip": (
                            "MiniMax H3 MODEL to patch for per-row mixed-timestep generation masking. "
                            "Connect this MODEL output to your guider."
                        )
                    },
                )
            }
        }

    RETURN_TYPES = ("MODEL",)
    RETURN_NAMES = ("patched_model",)
    FUNCTION = "patch"
    CATEGORY = "model/patch/minimax"
    DESCRIPTION = (
        "Temporarily enables AbleJones-style mixed-timestep masking for MiniMax H3 without "
        "editing ComfyUI core files. Only the returned MODEL is affected."
    )

    def patch(self, model):
        return (install_model_wrappers(model),)


class MiniMaxH3SetGenerationMask:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "av_latent": (
                    "LATENT",
                    {
                        "tooltip": (
                            "Joint source video+audio latent. Encode the source media, then combine it with "
                            "Concat AV Latent before connecting it here."
                        )
                    },
                ),
                "mask": (
                    "MASK",
                    {
                        "tooltip": (
                            "Video-space generation mask. It may contain one mask frame or a batch matching "
                            "the source video. The sampler resizes and snaps it to H3's patch grid."
                        )
                    },
                ),
                "mask_meaning": (
                    ["white = generate", "white = preserve"],
                    {"tooltip": "Choose whether white pixels are edited or protected."},
                ),
                "audio_mode": (
                    ["generate audio", "preserve source audio"],
                    {
                        "tooltip": (
                            "Generate audio normally, or pin the complete source-audio latent while editing video."
                        )
                    },
                ),
            }
        }

    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("masked_av_latent",)
    FUNCTION = "set_mask"
    CATEGORY = "latent/mask/minimax"
    DESCRIPTION = (
        "Attaches a MiniMax H3 video generation mask to a joint source AV latent. "
        "White normally means generate; black means preserve."
    )

    def set_mask(self, av_latent, mask, mask_meaning, audio_mode):
        import comfy.nested_tensor

        samples = av_latent.get("samples")
        if samples is None or not getattr(samples, "is_nested", False):
            raise ValueError(
                "H3 Set Generation Mask needs a joint video+audio LATENT. "
                "Use Concat AV Latent before this node."
            )
        streams = samples.unbind()
        if len(streams) < 2:
            raise ValueError("The connected AV latent does not contain both video and audio streams.")

        video_mask = mask.to(torch.float32).clamp(0.0, 1.0)
        if mask_meaning == "white = preserve":
            video_mask = 1.0 - video_mask
        audio_mask = (
            torch.ones_like(streams[1], dtype=torch.float32)
            if audio_mode == "generate audio"
            else torch.zeros_like(streams[1], dtype=torch.float32)
        )

        output = av_latent.copy()
        output["noise_mask"] = comfy.nested_tensor.NestedTensor((video_mask, audio_mask))
        return (output,)


class MiniMaxH3MaskGridPreview:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": (
                    "IMAGE",
                    {
                        "tooltip": (
                            "The exact resized H3 canvas frames that will be VAE-encoded. Width and height "
                            "must be divisible by 32 so the displayed cells align with H3."
                        )
                    },
                ),
                "mask": (
                    "MASK",
                    {
                        "tooltip": (
                            "Connect SAM3 Detect, SAM3 Track to Mask, a painted mask, or any standard "
                            "ComfyUI MASK. White is treated as the requested generation area."
                        )
                    },
                ),
                "cell_selection": (
                    list(GRID_MODES),
                    {
                        "tooltip": (
                            "Runtime exact reproduces H3's latent resize and per-row max decision. Other "
                            "modes select cells from direct pixel coverage and can be more or less conservative."
                        )
                    },
                ),
                "cell_adjust": (
                    "INT",
                    {
                        "default": 0,
                        "min": -8,
                        "max": 8,
                        "step": 1,
                        "tooltip": (
                            "Grow with positive values or shrink with negative values, measured in complete "
                            "32x32 H3 cells."
                        ),
                    },
                ),
                "overlay_opacity": (
                    "FLOAT",
                    {
                        "default": 0.38,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Opacity of the orange generated-cell overlay.",
                    },
                ),
                "show_grid": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Draw cyan boundaries for every 32x32 H3 cell.",
                    },
                ),
                "show_source_outline": (
                    "BOOLEAN",
                    {
                        "default": True,
                        "tooltip": "Draw the original SAM3 or input-mask boundary in white.",
                    },
                ),
            }
        }

    RETURN_TYPES = ("MASK", "IMAGE", "STRING")
    RETURN_NAMES = ("snapped_mask", "grid_preview", "grid_info")
    FUNCTION = "preview"
    CATEGORY = "latent/mask/minimax"
    DESCRIPTION = (
        "Shows exactly which 32x32 MiniMax H3 cells a SAM3 or standard ComfyUI mask will edit, "
        "and outputs the cell-snapped mask for H3 Set Generation Mask."
    )

    def preview(
        self,
        image,
        mask,
        cell_selection,
        cell_adjust,
        overlay_opacity,
        show_grid,
        show_source_outline,
    ):
        import torch.nn.functional as functional

        if image.ndim != 4 or image.shape[-1] < 3:
            raise ValueError(f"Expected ComfyUI IMAGE [frames,H,W,C], got {list(image.shape)}")
        frames, height, width = image.shape[:3]
        raw, snapped, cells = quantize_h3_pixel_mask(
            mask,
            frames,
            height,
            width,
            mode=cell_selection,
            cell_adjust=cell_adjust,
        )

        preview = image[..., :3].to(torch.float32).clamp(0.0, 1.0).clone()
        selected = snapped.to(device=preview.device, dtype=preview.dtype).unsqueeze(-1)
        orange = torch.tensor([1.0, 0.18, 0.0], device=preview.device, dtype=preview.dtype)
        opacity = float(overlay_opacity)
        preview = preview * (1.0 - selected * opacity) + orange * selected * opacity

        if show_grid:
            grid = torch.zeros((height, width), device=preview.device, dtype=torch.bool)
            grid[::H3_PIXEL_CELL, :] = True
            grid[:, ::H3_PIXEL_CELL] = True
            grid[-1, :] = True
            grid[:, -1] = True
            cyan = torch.tensor([0.0, 0.82, 1.0], device=preview.device, dtype=preview.dtype)
            grid_alpha = 0.72
            preview[:, grid] = preview[:, grid] * (1.0 - grid_alpha) + cyan * grid_alpha

        if show_source_outline:
            source = (raw.to(device=preview.device) >= 0.5).to(preview.dtype)
            dilated = functional.max_pool2d(source, kernel_size=3, stride=1, padding=1)
            eroded = 1.0 - functional.max_pool2d(1.0 - source, kernel_size=3, stride=1, padding=1)
            outline = ((dilated - eroded) > 0).movedim(1, -1)
            preview = torch.where(outline, torch.ones_like(preview), preview)

        counts = cells.sum(dim=(1, 2, 3)).to(torch.int64).cpu().tolist()
        cells_per_frame = cells.shape[-2] * cells.shape[-1]
        if len(counts) <= 8:
            count_text = ", ".join(str(value) for value in counts)
        else:
            count_text = f"min {min(counts)}, max {max(counts)}, total {sum(counts)}"
        grid_info = (
            f"H3 grid: {cells.shape[-1]}x{cells.shape[-2]} cells at {H3_PIXEL_CELL}x{H3_PIXEL_CELL} "
            f"pixels; selected cells per frame: {count_text} / {cells_per_frame}; "
            f"mode: {cell_selection}; cell adjustment: {cell_adjust:+d}."
        )
        return (snapped, preview, grid_info)


NODE_CLASS_MAPPINGS = {
    "MiniMaxH3TrimSourceAV": MiniMaxH3TrimSourceAV,
    "MiniMaxH3PerRowMaskPatch": MiniMaxH3PerRowMaskPatch,
    "MiniMaxH3SetGenerationMask": MiniMaxH3SetGenerationMask,
    "MiniMaxH3MaskGridPreview": MiniMaxH3MaskGridPreview,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MiniMaxH3TrimSourceAV": "MiniMax H3 Trim Source AV to 17k+5",
    "MiniMaxH3PerRowMaskPatch": "MiniMax H3 Per-Row Mask Patch (Experimental)",
    "MiniMaxH3SetGenerationMask": "MiniMax H3 Set Generation Mask",
    "MiniMaxH3MaskGridPreview": "MiniMax H3 Mask Grid Preview & Snap",
}