"""Save and load MiniMax H3 nested AV latents in ComfyUI.

MiniMax H3 stores its latent as::

    {"samples": NestedTensor((video_latent, audio_latent))}

The built-in ComfyUI Save Latent node expects a single tensor, so this module
stores the two H3 tensors under separate safetensors keys and reconstructs the
NestedTensor when loading.

File format (safetensors container, custom .h3latent extension):

    h3_video                 [B, 24, T_video, H, W]
    h3_audio                 [B, 32, 2, T_audio]
    h3_format_version        scalar/int marker
    latent_noise_mask        optional top-level LATENT noise_mask

Safetensors preserves the tensor dtype exactly. Tensors are copied to CPU and
made contiguous before saving.
"""

from __future__ import annotations

import hashlib
import json
import os
from typing import Any, Dict, Optional, Tuple

import torch
import safetensors.torch

import comfy.nested_tensor
import comfy.utils
import folder_paths


FORMAT_NAME = "comfyui_minimax_h3_nested_latent"
FORMAT_VERSION = 1
FILE_EXTENSION = ".h3latent"


def _unwrap_h3_samples(samples: Dict[str, Any]) -> Tuple[torch.Tensor, torch.Tensor]:
    """Extract and validate the video/audio tensors from a ComfyUI LATENT."""
    if not isinstance(samples, dict):
        raise TypeError("samples must be a ComfyUI LATENT dictionary")

    nested = samples.get("samples")
    if nested is None:
        raise ValueError("LATENT has no 'samples' entry")

    if isinstance(nested, comfy.nested_tensor.NestedTensor):
        tensors = list(nested.unbind())
    elif isinstance(nested, (tuple, list)):
        tensors = list(nested)
    else:
        raise TypeError(
            "Expected MiniMax H3 NestedTensor(video, audio), got "
            f"{type(nested).__name__}"
        )

    if len(tensors) != 2:
        raise ValueError(
            "MiniMax H3 LATENT must contain exactly two tensors "
            f"(video, audio); got {len(tensors)}"
        )

    video, audio = tensors
    if not torch.is_tensor(video) or not torch.is_tensor(audio):
        raise TypeError("Both H3 nested objects must be torch tensors")

    if video.ndim != 5 or video.shape[1] != 24:
        raise ValueError(
            "Expected H3 video latent [B, 24, T, H, W], got "
            f"shape={tuple(video.shape)}"
        )

    if audio.ndim != 4 or audio.shape[1] != 32 or audio.shape[2] != 2:
        raise ValueError(
            "Expected H3 audio latent [B, 32, 2, T], got "
            f"shape={tuple(audio.shape)}"
        )

    if video.shape[0] != audio.shape[0]:
        raise ValueError(
            "H3 video/audio batch sizes differ: "
            f"video={video.shape[0]}, audio={audio.shape[0]}"
        )

    return video, audio


def _tensor_for_save(tensor: torch.Tensor) -> torch.Tensor:
    """Safetensors requires dense, contiguous tensors."""
    if tensor.layout != torch.strided:
        tensor = tensor.to_dense()
    return tensor.detach().to(device="cpu").contiguous()


def _json_metadata(samples: Dict[str, Any]) -> str:
    """Preserve small, JSON-safe LATENT fields other than tensor payloads."""
    metadata: Dict[str, Any] = {}
    for key, value in samples.items():
        if key in {"samples", "noise_mask"}:
            continue
        try:
            json.dumps(value)
        except (TypeError, ValueError):
            continue
        metadata[key] = value
    return json.dumps(metadata, ensure_ascii=False, separators=(",", ":"))


def _safe_join(root: str, relative_path: str) -> str:
    """Resolve a path below root and reject directory traversal."""
    if not isinstance(relative_path, str) or not relative_path.strip():
        raise ValueError("latent file path is empty")

    relative_path = relative_path.strip().replace("\\", os.sep).replace("/", os.sep)
    if os.path.isabs(relative_path):
        raise ValueError("Use a path relative to ComfyUI input/output, not an absolute path")

    root_abs = os.path.abspath(root)
    full = os.path.abspath(os.path.join(root_abs, os.path.normpath(relative_path)))
    try:
        inside = os.path.commonpath((root_abs, full)) == root_abs
    except ValueError:
        inside = False
    if not inside:
        raise ValueError("latent file path escapes the selected ComfyUI directory")
    return full


def _parse_file_reference(file_ref: str, location: str) -> Tuple[str, str, str]:
    """Return (absolute path, effective location, relative path).

    Save returns references such as ``output:h3_latents/foo.h3latent``.
    The loader also accepts a plain relative path plus the location widget.
    """
    raw = (file_ref or "").strip()
    effective = location

    for prefix in ("output:", "input:"):
        if raw.lower().startswith(prefix):
            effective = prefix[:-1]
            raw = raw[len(prefix):]
            break

    if effective == "auto":
        effective = "output"

    if effective == "output":
        root = folder_paths.get_output_directory()
    elif effective == "input":
        root = folder_paths.get_input_directory()
    else:
        raise ValueError(f"Unsupported location: {effective}")

    full = _safe_join(root, raw)
    return full, effective, raw.replace(os.sep, "/")


def _human_bytes(value: int) -> str:
    size = float(value)
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if size < 1024.0 or unit == "TiB":
            return f"{size:.2f} {unit}"
        size /= 1024.0
    return f"{value} B"


class SaveMiniMaxH3NestedLatent:
    """Save both tensors from an H3 NestedTensor into one safetensors file."""

    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "samples": ("LATENT",),
                "filename_prefix": (
                    "STRING",
                    {"default": "h3_latents/ComfyUI"},
                ),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
            },
        }

    RETURN_TYPES = ("LATENT", "STRING", "STRING")
    RETURN_NAMES = ("samples", "saved_file", "info")
    FUNCTION = "save"
    OUTPUT_NODE = True
    CATEGORY = "MiniMax H3/Experimental"
    DESCRIPTION = (
        "Saves the H3 NestedTensor video and audio latents losslessly into one "
        ".h3latent safetensors file."
    )

    def save(
        self,
        samples,
        filename_prefix="h3_latents/ComfyUI",
        prompt=None,
        extra_pnginfo=None,
    ):
        video, audio = _unwrap_h3_samples(samples)

        full_output_folder, filename, counter, subfolder, _ = (
            folder_paths.get_save_image_path(filename_prefix, self.output_dir)
        )
        os.makedirs(full_output_folder, exist_ok=True)

        file_name = f"{filename}_{counter:05}_{FILE_EXTENSION}"
        absolute_path = os.path.join(full_output_folder, file_name)

        tensors: Dict[str, torch.Tensor] = {
            "h3_video": _tensor_for_save(video),
            "h3_audio": _tensor_for_save(audio),
            "h3_format_version": torch.tensor([FORMAT_VERSION], dtype=torch.int32),
        }

        noise_mask = samples.get("noise_mask")
        if torch.is_tensor(noise_mask):
            tensors["latent_noise_mask"] = _tensor_for_save(noise_mask)

        metadata: Dict[str, str] = {
            "format": FORMAT_NAME,
            "version": str(FORMAT_VERSION),
            "video_shape": json.dumps(list(video.shape)),
            "audio_shape": json.dumps(list(audio.shape)),
            "video_dtype": str(video.dtype),
            "audio_dtype": str(audio.dtype),
            "latent_metadata": _json_metadata(samples),
        }

        if prompt is not None:
            metadata["prompt"] = json.dumps(prompt, ensure_ascii=False, default=str)
        if extra_pnginfo is not None:
            metadata["extra_pnginfo"] = json.dumps(
                extra_pnginfo,
                ensure_ascii=False,
                default=str,
            )

        # ComfyUI's helper writes a safetensors container regardless of our
        # custom extension. Every tensor has already been made CPU-contiguous.
        comfy.utils.save_torch_file(tensors, absolute_path, metadata=metadata)

        relative = os.path.relpath(absolute_path, self.output_dir).replace(os.sep, "/")
        file_ref = f"output:{relative}"
        file_size = os.path.getsize(absolute_path)
        info = (
            f"Saved {file_ref}; video={tuple(video.shape)} {video.dtype}; "
            f"audio={tuple(audio.shape)} {audio.dtype}; size={_human_bytes(file_size)}"
        )

        return {
            "ui": {"text": [info]},
            "result": (samples, file_ref, info),
        }


class LoadMiniMaxH3NestedLatent:
    """Load an H3 AV latent and reconstruct NestedTensor((video, audio))."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "file": (
                    "STRING",
                    {"default": "output:h3_latents/ComfyUI_00001_.h3latent"},
                ),
                "location": (
                    ["auto", "output", "input"],
                    {"default": "auto"},
                ),
            }
        }

    RETURN_TYPES = ("LATENT", "STRING")
    RETURN_NAMES = ("samples", "info")
    FUNCTION = "load"
    CATEGORY = "MiniMax H3/Experimental"
    DESCRIPTION = (
        "Loads a .h3latent file and rebuilds the MiniMax H3 "
        "NestedTensor(video, audio) LATENT."
    )

    def load(self, file, location="auto"):
        absolute_path, effective_location, relative = _parse_file_reference(
            file,
            location,
        )
        if not os.path.isfile(absolute_path):
            raise FileNotFoundError(
                f"H3 latent file does not exist: {effective_location}:{relative}"
            )

        tensors = safetensors.torch.load_file(absolute_path, device="cpu")

        required = {"h3_video", "h3_audio", "h3_format_version"}
        missing = sorted(required.difference(tensors))
        if missing:
            raise ValueError(
                "Not a supported H3 nested latent file; missing keys: "
                + ", ".join(missing)
            )

        version_tensor = tensors["h3_format_version"]
        version = int(version_tensor.reshape(-1)[0].item())
        if version != FORMAT_VERSION:
            raise ValueError(
                f"Unsupported H3 latent format version {version}; "
                f"this node supports {FORMAT_VERSION}"
            )

        video = tensors["h3_video"]
        audio = tensors["h3_audio"]

        # Validate using the same rules as saving without requiring a temporary
        # NestedTensor on GPU.
        reconstructed = {
            "samples": comfy.nested_tensor.NestedTensor((video, audio))
        }
        _unwrap_h3_samples(reconstructed)

        if "latent_noise_mask" in tensors:
            reconstructed["noise_mask"] = tensors["latent_noise_mask"]

        # Restore small JSON-safe top-level LATENT fields when present.
        try:
            with safetensors.safe_open(absolute_path, framework="pt", device="cpu") as handle:
                file_metadata = handle.metadata() or {}
            latent_metadata = json.loads(file_metadata.get("latent_metadata", "{}"))
            if isinstance(latent_metadata, dict):
                reconstructed.update(latent_metadata)
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            # The tensors are the load-bearing data; metadata failure should not
            # invalidate an otherwise usable latent.
            pass

        file_size = os.path.getsize(absolute_path)
        finite_video = bool(torch.isfinite(video).all().item())
        finite_audio = bool(torch.isfinite(audio).all().item())
        info = (
            f"Loaded {effective_location}:{relative}; "
            f"video={tuple(video.shape)} {video.dtype}; "
            f"audio={tuple(audio.shape)} {audio.dtype}; "
            f"finite(video/audio)={finite_video}/{finite_audio}; "
            f"size={_human_bytes(file_size)}"
        )
        return reconstructed, info

    @classmethod
    def IS_CHANGED(cls, file, location="auto"):
        try:
            absolute_path, _, _ = _parse_file_reference(file, location)
            digest = hashlib.sha256()
            with open(absolute_path, "rb") as handle:
                for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                    digest.update(chunk)
            return digest.hexdigest()
        except (OSError, ValueError):
            return float("nan")

    @classmethod
    def VALIDATE_INPUTS(cls, file, location="auto"):
        try:
            absolute_path, effective_location, relative = _parse_file_reference(
                file,
                location,
            )
        except ValueError as exc:
            return str(exc)
        if not os.path.isfile(absolute_path):
            return f"Invalid H3 latent file: {effective_location}:{relative}"
        if not relative.lower().endswith(FILE_EXTENSION):
            return f"Expected a {FILE_EXTENSION} file"
        return True