"""Benchmark MiniMax H3 video VAE decode variants on a saved latent."""

from __future__ import annotations

import argparse
import json
import subprocess
import time
from pathlib import Path

import torch


def sync() -> None:
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def timed(stage: str, rows: list[dict], fn):
    sync()
    start = time.perf_counter()
    value = fn()
    sync()
    elapsed = time.perf_counter() - start
    rows.append({"stage": stage, "seconds": elapsed})
    return value


def dtype_from_name(name: str) -> torch.dtype:
    return {
        "float32": torch.float32,
        "float16": torch.float16,
        "bfloat16": torch.bfloat16,
    }[name]


def ffmpeg_command(loglevel: str, *parts: str) -> list[str]:
    return ["ffmpeg", "-hide_banner", "-loglevel", loglevel, *parts]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--latent", type=Path, required=True)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--metrics", type=Path, required=True)
    parser.add_argument("--implementation", choices=("direct", "upstream"), default="direct")
    parser.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="float32")
    parser.add_argument("--no-tiling", action="store_true")
    parser.add_argument("--tile-size", type=int)
    parser.add_argument("--tile-overlap", type=int)
    parser.add_argument("--trace-calls", action="store_true")
    parser.add_argument("--skip-video", action="store_true")
    parser.add_argument("--ffmpeg-loglevel", default="error")
    args = parser.parse_args()

    rows: list[dict] = []
    call_rows: list[dict] = []
    dtype = dtype_from_name(args.dtype)

    state = timed("latent_load", rows, lambda: 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":
        from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE

        vae = timed(
            "vae_load",
            rows,
            lambda: MiniMaxH3VideoVAE.from_safetensors(
                "/vae/minimax_h3_video_vae_fp16.safetensors",
                device="cuda",
                tiling=not args.no_tiling,
                dtype=dtype,
            ).eval(),
        )
    else:
        from safetensors.torch import load_file

        from h3_blackwell_runtime.upstream_vae import MiniMaxH3VideoVAE

        def load_upstream():
            model = MiniMaxH3VideoVAE(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) for key, value in checkpoint.items()}
            model.load_state_dict(checkpoint, strict=True)
            return model

        vae = timed("vae_load", rows, load_upstream)

    if args.tile_size is not None:
        vae.tile_size = args.tile_size
    if args.tile_overlap is not None:
        vae.tile_overlap_min = args.tile_overlap

    if args.trace_calls:
        original_decode_pixels = vae._decode_pixels

        def traced_decode_pixels(z):
            sync()
            start = time.perf_counter()
            result = original_decode_pixels(z)
            sync()
            call_rows.append({
                "index": len(call_rows),
                "latent_shape": tuple(z.shape),
                "seconds": time.perf_counter() - start,
            })
            return result

        vae._decode_pixels = traced_decode_pixels

    with torch.inference_mode():
        decoded = timed("vae_decode", rows, lambda: vae.decode(latent.to(dtype))[:, :, :frames])
        pixels = timed(
            "pixelize_cpu",
            rows,
            lambda: ((decoded[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu(),
        )

    if args.output is not None and not args.skip_video:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        raw = args.output.with_suffix(".rgb")
        timed("raw_write", rows, lambda: pixels.numpy().tofile(raw))
        timed(
            "ffmpeg_encode",
            rows,
            lambda: subprocess.run(
                ffmpeg_command(
                    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()

    metrics = {
        "implementation": args.implementation,
        "dtype": args.dtype,
        "tiling": not args.no_tiling,
        "tile_size": vae.tile_size,
        "tile_overlap_min": vae.tile_overlap_min,
        "frames": frames,
        "latent_shape": tuple(latent.shape),
        "output": str(args.output) if args.output is not None else None,
        "stages": rows,
        "decode_pixel_calls": call_rows,
        "total_seconds": sum(row["seconds"] for row in rows),
    }
    args.metrics.parent.mkdir(parents=True, exist_ok=True)
    args.metrics.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(json.dumps(metrics, indent=2), flush=True)


if __name__ == "__main__":
    main()