"""
H3 Bundle — sauvegarde et rechargement de latents AV + conditioning MiniMax H3.

Objectif : découpler la passe de génération (locale) de la passe d'upscaling
(instance Vast.ai), sans avoir à réimporter les images de référence ni le
CLIP qwen3vl 32B.

Installation :
    cp -r h3_bundle /workspace/ComfyUI/custom_nodes/
    (idem en local dans ComfyUI/custom_nodes/)
"""

import os
import json
import time
import shutil
import torch

import folder_paths


# --------------------------------------------------------------------------
# helpers
# --------------------------------------------------------------------------

def _to_cpu(obj):
    """Descend récursivement la structure et rapatrie tous les tenseurs en CPU."""
    if torch.is_tensor(obj):
        return obj.detach().cpu()
    if isinstance(obj, dict):
        return {k: _to_cpu(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return type(obj)(_to_cpu(v) for v in obj)
    return obj


def _shapes(obj):
    """Cartographie des shapes/dtypes — c'est ce qui te sert à câbler l'amont."""
    if torch.is_tensor(obj):
        return f"Tensor{tuple(obj.shape)} {str(obj.dtype).replace('torch.', '')}"
    if isinstance(obj, dict):
        return {k: _shapes(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_shapes(v) for v in obj]
    if isinstance(obj, (int, float, str, bool, type(None))):
        return repr(obj)
    return type(obj).__name__


def _dir_size(path):
    total = 0
    for root, _, files in os.walk(path):
        for f in files:
            total += os.path.getsize(os.path.join(root, f))
    return total


def _human(n):
    for unit in ("B", "KB", "MB", "GB"):
        if n < 1024:
            return f"{n:.1f} {unit}"
        n /= 1024
    return f"{n:.1f} TB"


def _bundles_root():
    return os.path.join(folder_paths.get_output_directory(), "h3_bundles")


# --------------------------------------------------------------------------
# preview
# --------------------------------------------------------------------------

def _tensor_to_pil(frame, max_width=512):
    """[H,W,C] float 0-1 -> PIL, redimensionné pour rester léger."""
    from PIL import Image
    import numpy as np
    arr = (frame.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
    img = Image.fromarray(arr)
    if img.width > max_width:
        ratio = max_width / img.width
        img = img.resize((max_width, int(img.height * ratio)), Image.LANCZOS)
    return img


def _sample_frames(images, count):
    """Échantillonne count frames réparties sur toute la séquence."""
    total = images.shape[0]
    if total <= count:
        return list(range(total))
    step = (total - 1) / (count - 1) if count > 1 else 1
    return [min(total - 1, int(round(i * step))) for i in range(count)]


def _write_wav(audio, path):
    """AUDIO ComfyUI ({waveform:[B,C,S], sample_rate}) -> wav 16 bits."""
    import wave
    import numpy as np
    wav = audio["waveform"]
    if wav.dim() == 3:
        wav = wav[0]
    data = wav.transpose(0, 1).cpu().numpy()
    data = (data.clip(-1, 1) * 32767).astype(np.int16)
    with wave.open(path, "wb") as fh:
        fh.setnchannels(data.shape[1])
        fh.setsampwidth(2)
        fh.setframerate(int(audio["sample_rate"]))
        fh.writeframes(data.tobytes())


def _encode_mp4(images, fps, out_path, audio=None):
    """Encode la séquence en mp4 via ffmpeg. Retourne True si ça a marché."""
    import subprocess
    import tempfile

    if shutil.which("ffmpeg") is None:
        print("[SaveH3Bundle] ffmpeg introuvable, preview mp4 ignorée.")
        return False

    tmp = tempfile.mkdtemp(prefix="h3prev_")
    try:
        for i in range(images.shape[0]):
            _tensor_to_pil(images[i], max_width=720).save(
                os.path.join(tmp, f"{i:06d}.png"))

        cmd = ["ffmpeg", "-y", "-loglevel", "error",
               "-framerate", str(fps), "-i", os.path.join(tmp, "%06d.png")]

        wav_path = None
        if audio is not None:
            wav_path = os.path.join(tmp, "audio.wav")
            try:
                _write_wav(audio, wav_path)
                cmd += ["-i", wav_path]
            except Exception as exc:
                print(f"[SaveH3Bundle] audio preview ignoré : {exc}")
                wav_path = None

        cmd += ["-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "20"]
        if wav_path:
            cmd += ["-c:a", "aac", "-shortest"]
        cmd.append(out_path)

        subprocess.run(cmd, check=True)
        return True
    except Exception as exc:
        print(f"[SaveH3Bundle] échec encodage preview : {exc}")
        return False
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


# --------------------------------------------------------------------------
# SAVE
# --------------------------------------------------------------------------

class SaveH3Bundle:
    """
    Regroupe tout ce dont la passe d'upscale a besoin dans un dossier unique.

    Le latent MiniMax H3 porte déjà vidéo + audio dans sa structure nested,
    donc un seul torch.save suffit — contrairement au SaveLatent natif qui
    ne sérialise que la clé `samples`.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "latent": ("LATENT",),
                "bundle_name": ("STRING", {"default": "run"}),
                "mode": (["save", "replace_last", "skip"],),
            },
            "optional": {
                "conditioning": ("CONDITIONING",),
                "video": ("VIDEO",),
                "video_path": ("STRING", {"default": ""}),
                "images": ("IMAGE",),
                "audio": ("AUDIO",),
                "preview_frames": ("INT", {"default": 6, "min": 0, "max": 24}),
                "preview_mp4": ("BOOLEAN", {"default": True}),
                "keep_last_n": ("INT", {"default": 0, "min": 0, "max": 999}),
                "width": ("INT", {"default": 0, "min": 0, "max": 16384}),
                "height": ("INT", {"default": 0, "min": 0, "max": 16384}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
                "fps": ("INT", {"default": 24, "min": 1, "max": 240}),
                "generation_seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 600.0}),
                "notes": ("STRING", {"default": "", "multiline": True}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("bundle_path",)
    FUNCTION = "save"
    OUTPUT_NODE = True
    CATEGORY = "vast/h3"

    @staticmethod
    def _existing_bundles():
        """Bundles valides, du plus ancien au plus récent."""
        root = _bundles_root()
        if not os.path.isdir(root):
            return []
        found = []
        for name in sorted(os.listdir(root)):
            path = os.path.join(root, name)
            if os.path.isdir(path) and os.path.isfile(os.path.join(path, "latent.pt")):
                found.append(path)
        return found

    @staticmethod
    def _delete_bundle(path):
        """Suppression volontairement paranoïaque : jamais hors de h3_bundles,
        et uniquement un dossier qui contient bien un latent.pt."""
        root = os.path.abspath(_bundles_root())
        target = os.path.abspath(path)
        if not target.startswith(root + os.sep):
            print(f"[SaveH3Bundle] REFUS suppression hors bundles : {target}")
            return False
        if not os.path.isfile(os.path.join(target, "latent.pt")):
            print(f"[SaveH3Bundle] REFUS suppression, pas un bundle : {target}")
            return False
        shutil.rmtree(target)
        print(f"[SaveH3Bundle] supprimé : {os.path.basename(target)}")
        return True

    def save(self, latent, bundle_name, mode="save", conditioning=None,
             video=None, video_path="", images=None, audio=None,
             preview_frames=6, preview_mp4=True,
             keep_last_n=0, width=0, height=0, seed=0, fps=24,
             generation_seconds=0.0, notes=""):

        if mode == "skip":
            print("[SaveH3Bundle] mode 'skip' — rien n'est écrit sur le disque.")
            return {"ui": {"text": ["skip"]}, "result": ("",)}

        if mode == "replace_last":
            existing = self._existing_bundles()
            if existing:
                self._delete_bundle(existing[-1])
            else:
                print("[SaveH3Bundle] 'replace_last' : aucun bundle à supprimer.")

        safe = "".join(c for c in bundle_name if c.isalnum() or c in "-_") or "run"
        stamp = time.strftime("%Y%m%d_%H%M%S")
        root = os.path.join(_bundles_root(), f"{safe}_{stamp}")
        os.makedirs(root, exist_ok=True)

        manifest = {
            "bundle_name": safe,
            "created": stamp,
            "width": width,
            "height": height,
            "seed": seed,
            "fps": fps,
            "generation_seconds": generation_seconds,
            "notes": notes,
            "has_conditioning": conditioning is not None,
        }

        torch.save(_to_cpu(latent), os.path.join(root, "latent.pt"))
        manifest["latent_structure"] = _shapes(latent)

        if conditioning is not None:
            torch.save(_to_cpu(conditioning), os.path.join(root, "cond.pt"))
            manifest["cond_structure"] = _shapes(conditioning)

        manifest["files"] = {
            f: _human(os.path.getsize(os.path.join(root, f)))
            for f in sorted(os.listdir(root))
        }
        manifest["total_size"] = _human(_dir_size(root))

        with open(os.path.join(root, "manifest.json"), "w") as fh:
            json.dump(manifest, fh, indent=2, default=str)

        print("\n" + "=" * 66)
        print(f"[SaveH3Bundle] {root}")
        print(f"  total : {manifest['total_size']}")
        print("  structure latent :")
        print(json.dumps(manifest["latent_structure"], indent=4, default=str))
        if conditioning is not None:
            print("  structure conditioning :")
            print(json.dumps(manifest["cond_structure"], indent=4, default=str))
        print("=" * 66 + "\n")

        # ---- preview ----------------------------------------------------
        ui_images = []
        dest_mp4 = os.path.join(root, "preview.mp4")
        got_video = False

        # 1. objet VIDEO natif (sortie de CreateVideo) — écriture directe
        if video is not None:
            try:
                video.save_to(dest_mp4)
                got_video = True
                print(f"[SaveH3Bundle] vidéo écrite : {dest_mp4}")
            except Exception as exc:
                print(f"[SaveH3Bundle] save_to a échoué ({exc}), on tente autre chose.")

        # 2. chemin d'un fichier déjà écrit par un autre node — simple copie
        if not got_video and video_path:
            src = video_path
            if not os.path.isabs(src):
                src = os.path.join(folder_paths.get_output_directory(), src)
            if os.path.isfile(src):
                shutil.copy2(src, os.path.join(root, os.path.basename(src)))
                got_video = True
                print(f"[SaveH3Bundle] vidéo copiée : {os.path.basename(src)}")
            else:
                print(f"[SaveH3Bundle] video_path introuvable : {src}")

        # 3. dernier recours — ré-encodage depuis les frames
        if not got_video and preview_mp4 and images is not None and images.shape[0] > 0:
            if _encode_mp4(images, fps, dest_mp4, audio=audio):
                got_video = True

        manifest["has_video"] = got_video

        if images is not None and images.shape[0] > 0 and preview_frames > 0:
            sub = os.path.join("h3_bundles", os.path.basename(root))
            for n, idx in enumerate(_sample_frames(images, preview_frames)):
                fname = f"preview_{n:02d}.png"
                _tensor_to_pil(images[idx]).save(os.path.join(root, fname))
                ui_images.append({
                    "filename": fname,
                    "subfolder": sub,
                    "type": "output",
                })

        manifest["files"] = {
            f: _human(os.path.getsize(os.path.join(root, f)))
            for f in sorted(os.listdir(root)) if os.path.isfile(os.path.join(root, f))
        }
        manifest["total_size"] = _human(_dir_size(root))
        with open(os.path.join(root, "manifest.json"), "w") as fh:
            json.dump(manifest, fh, indent=2, default=str)

        if keep_last_n > 0:
            everything = self._existing_bundles()
            for old in everything[:-keep_last_n]:
                self._delete_bundle(old)

        return {"ui": {"images": ui_images, "text": [root]}, "result": (root,)}


# --------------------------------------------------------------------------
# LOAD
# --------------------------------------------------------------------------

class LoadH3BundleBatch:
    """
    Recharge un bundle. ComfyUI n'a pas de boucle interne : le batch séquentiel
    se fait en empilant N runs dans la queue.

    mode = "index"        -> tu choisis le bundle à la main (debug / rerun)
    mode = "next_pending" -> avance tout seul dans le dossier, un bundle par
                             run. Tu cliques Queue N fois (ou tu mets le batch
                             count à N) et il déroule la liste.

    Un bundle est considéré comme traité si un fichier de sortie portant son
    nom existe déjà dans output/<output_subdir>/. Pas de marqueur : un run qui
    plante est repris automatiquement au run suivant.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "bundles_dir": ("STRING", {"default": "h3_bundles"}),
                "mode": (["next_pending", "index"],),
                "index": ("INT", {"default": 0, "min": 0, "max": 9999}),
            },
            "optional": {
                "output_subdir": ("STRING", {"default": "minimax-h3"}),
            },
        }

    RETURN_TYPES = ("LATENT", "CONDITIONING", "STRING", "INT", "INT", "INT")
    RETURN_NAMES = ("latent", "conditioning", "filename_prefix",
                    "width", "height", "remaining")
    FUNCTION = "load"
    CATEGORY = "vast/h3"

    @classmethod
    def IS_CHANGED(cls, mode, **kwargs):
        # en next_pending l'état change à chaque run : on casse le cache
        if mode == "next_pending":
            return float("nan")
        return json.dumps(kwargs, sort_keys=True, default=str)

    def _resolve_dir(self, bundles_dir):
        if os.path.isabs(bundles_dir):
            return bundles_dir
        for base in (folder_paths.get_output_directory(),
                     folder_paths.get_input_directory()):
            cand = os.path.join(base, bundles_dir)
            if os.path.isdir(cand):
                return cand
        return os.path.join(folder_paths.get_output_directory(), bundles_dir)

    def _list_bundles(self, root):
        if not os.path.isdir(root):
            raise FileNotFoundError(f"Dossier de bundles introuvable : {root}")
        out = []
        for name in sorted(os.listdir(root)):
            path = os.path.join(root, name)
            if os.path.isdir(path) and os.path.isfile(os.path.join(path, "latent.pt")):
                out.append(path)
        if not out:
            raise FileNotFoundError(f"Aucun bundle valide dans {root}")
        return out

    def _has_output(self, bundle_path, output_subdir):
        """Un bundle est traité si un fichier de sortie porte son nom."""
        out_dir = os.path.join(folder_paths.get_output_directory(), output_subdir)
        if not os.path.isdir(out_dir):
            return False
        name = os.path.basename(bundle_path)
        return any(f.startswith(name) for f in os.listdir(out_dir))

    def load(self, bundles_dir, mode, index, output_subdir="minimax-h3"):
        root = self._resolve_dir(bundles_dir)
        bundles = self._list_bundles(root)

        if mode == "index":
            if index >= len(bundles):
                raise IndexError(f"index {index} hors bornes ({len(bundles)} bundles)")
            target = bundles[index]
            remaining = len(bundles) - index - 1
        else:
            pending = [b for b in bundles if not self._has_output(b, output_subdir)]
            if not pending:
                raise RuntimeError(
                    f"Les {len(bundles)} bundles de {root} ont déjà une sortie dans "
                    f"output/{output_subdir}/. Supprime les mp4 correspondants pour "
                    f"les relancer, ou passe en mode 'index'."
                )
            target = pending[0]
            remaining = len(pending) - 1

        latent = torch.load(os.path.join(target, "latent.pt"),
                            map_location="cpu", weights_only=False)

        cond_file = os.path.join(target, "cond.pt")
        conditioning = []
        if os.path.isfile(cond_file):
            conditioning = torch.load(cond_file, map_location="cpu", weights_only=False)

        width = height = 0
        manifest_file = os.path.join(target, "manifest.json")
        if os.path.isfile(manifest_file):
            with open(manifest_file) as fh:
                m = json.load(fh)
            width = int(m.get("width", 0) or 0)
            height = int(m.get("height", 0) or 0)

        name = os.path.basename(target)
        prefix = f"minimax-h3/{name}_upscaled"

        print(f"[LoadH3BundleBatch] {name} | cond={'oui' if conditioning else 'non'} "
              f"| {width}x{height} | restants : {remaining}")

        return (latent, conditioning, prefix, width, height, remaining)


NODE_CLASS_MAPPINGS = {
    "SaveH3Bundle": SaveH3Bundle,
    "LoadH3BundleBatch": LoadH3BundleBatch,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SaveH3Bundle": "Save H3 Bundle (latent AV + cond)",
    "LoadH3BundleBatch": "Load H3 Bundle (batch séquentiel)",
}