"""Decode a saved H3 video latent with direct VAE settings."""

import argparse
import subprocess
from pathlib import Path

import torch

from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE, dtype_from_name


parser = argparse.ArgumentParser()
parser.add_argument("--latent", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--no-tiling", action="store_true")
parser.add_argument("--implementation", choices=("direct", "upstream"), default="direct")
parser.add_argument("--frames-dir", type=Path)
parser.add_argument("--ffmpeg-loglevel", default="error")
parser.add_argument("--quiet", action="store_true")
parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default="float16")
parser.add_argument("--vae-tile-size", type=int, default=256)
args = parser.parse_args()

state = torch.load(args.latent, map_location="cuda", weights_only=False)
latent = state["latent"].to("cuda") if isinstance(state, dict) else state.to("cuda")
frames = int(state.get("frames", latent.shape[2] * 4)) if isinstance(state, dict) else latent.shape[2] * 4

if args.implementation == "direct":
    vae = MiniMaxH3VideoVAE.from_safetensors(
        "/vae/minimax_h3_video_vae_fp16.safetensors",
        device="cuda",
        tiling=not args.no_tiling,
        dtype=dtype_from_name(args.vae_dtype),
    ).eval()
    vae.tile_size = args.vae_tile_size
else:
    from safetensors.torch import load_file

    from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE as UpstreamMiniMaxH3VideoVAE

    vae = UpstreamMiniMaxH3VideoVAE(tiling=not args.no_tiling).to("cuda").eval()
    checkpoint = load_file("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda")
    checkpoint = {key: value.to(dtype=dtype_from_name(args.vae_dtype)) for key, value in checkpoint.items()}
    missing, unexpected = vae.load_state_dict(checkpoint, strict=False)
    vae = vae.to(dtype=dtype_from_name(args.vae_dtype))
    vae.tile_size = args.vae_tile_size
    if not args.quiet:
        print({"upstream_missing": len(missing), "upstream_unexpected": len(unexpected)}, flush=True)
with torch.inference_mode():
    pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
pixels = ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()

if args.frames_dir is not None:
    from PIL import Image

    args.frames_dir.mkdir(parents=True, exist_ok=True)
    for index, frame in enumerate(pixels.numpy()):
        Image.fromarray(frame).save(args.frames_dir / f"frame_{index:04d}.png")

args.output.parent.mkdir(parents=True, exist_ok=True)
raw = args.output.with_suffix(".rgb")
pixels.numpy().tofile(raw)
subprocess.run([
    "ffmpeg", "-hide_banner", "-loglevel", args.ffmpeg_loglevel,
    "-y", "-f", "rawvideo", "-pixel_format", "rgb24",
    "-video_size", f"{pixels.shape[2]}x{pixels.shape[1]}", "-framerate", "24",
    "-i", str(raw), "-an", "-c:v", "libx264", "-pix_fmt", "yuv420p", str(args.output),
], check=True)
raw.unlink()
if not args.quiet:
    print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape), "tiling": not args.no_tiling, "implementation": args.implementation, "vae_dtype": args.vae_dtype, "vae_tile_size": vae.tile_size})