"""Generate a minimal direct, video-only H3 T2V preview without ComfyUI."""

import argparse
import os
from pathlib import Path
import subprocess
import time
from datetime import datetime, timezone
import warnings

warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning)

import torch

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
from h3_blackwell_runtime.packing import H3PromptPacker
from h3_blackwell_runtime.qwen3vl_text import Qwen3VLPromptConditioner
from h3_blackwell_runtime.sampler import sample_video_res_multistep
from h3_blackwell_runtime.t2v import random_av_latents
from h3_blackwell_runtime.token_refiner import H3TokenRefiner
from h3_blackwell_runtime.vae_decoder import MiniMaxH3VideoVAE


parser = argparse.ArgumentParser()
parser.add_argument("--prompt", default="A brass-and-paper dragon flies above a rain-washed old city at blue hour.")
parser.add_argument("--output", type=Path, default=Path("/output/direct-h3-preview.mp4"))
parser.add_argument("--width", type=int, default=320)
parser.add_argument("--height", type=int, default=192)
parser.add_argument("--frames", type=int, default=22)
parser.add_argument("--steps", type=int, default=12)
parser.add_argument("--seed", type=int, default=440204)
parser.add_argument("--attention", choices=("sage2", "sdpa", "sage3"), default="sage2")
parser.add_argument("--model-timesteps-capture", type=Path, help="Directory containing captured input_XX.pt H3 timesteps for strict parity checks.")
parser.add_argument("--profile-memory", action="store_true")
args = parser.parse_args()
started = time.perf_counter()
last_report = started


def report_memory(stage: str) -> None:
    global last_report
    if not args.profile_memory:
        return
    now = time.perf_counter()
    rss_kb = 0
    try:
        with open("/proc/self/status", encoding="utf-8") as file:
            for line in file:
                if line.startswith("VmRSS:"):
                    rss_kb = int(line.split()[1])
                    break
    except OSError:
        pass
    cuda_alloc = torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0.0
    cuda_reserved = torch.cuda.memory_reserved() / 1024**3 if torch.cuda.is_available() else 0.0
    print({"stage": stage, "ts": datetime.now(timezone.utc).isoformat(), "epoch_s": round(time.time(), 3), "elapsed_s": round(now - started, 3), "delta_s": round(now - last_report, 3), "rss_gb": round(rss_kb / 1024**2, 3), "cuda_alloc_gb": round(cuda_alloc, 3), "cuda_reserved_gb": round(cuda_reserved, 3), "fast_safetensors": os.getenv("H3_FAST_SAFETENSORS", ""), "disable_mmap": os.getenv("H3_DISABLE_MMAP", "")}, flush=True)
    last_report = now

report_memory("start")
checkpoint = H3Checkpoint("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
conditioner = Qwen3VLPromptConditioner(
    "/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors",
    "/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer",
)
report_memory("qwen_loaded")
video, audio, frames = random_av_latents(args.width, args.height, args.frames, args.seed)
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend=args.attention).eval()
report_memory("h3_loaded")
text = H3TokenRefiner(checkpoint, attention_backend=args.attention)(conditioner(args.prompt))
report_memory("text_conditioned")
model_timesteps = None
if args.model_timesteps_capture is not None:
    model_timesteps = [
        torch.load(args.model_timesteps_capture / f"input_{index:02d}.pt", map_location="cuda", weights_only=False)["timesteps"]
        for index in range(args.steps)
    ]
latent = sample_video_res_multistep(model, H3PromptPacker(checkpoint), text, video, audio, steps=args.steps, model_timesteps=model_timesteps)
report_memory("sampled")
vae = MiniMaxH3VideoVAE.from_safetensors("/vae/minimax_h3_video_vae_fp16.safetensors", device="cuda").eval()
report_memory("vae_loaded")
with torch.inference_mode():
    pixels = vae.decode(latent.to(next(vae.parameters()).dtype))[:, :, :frames]
report_memory("vae_decoded")
pixels = ((pixels[0].permute(1, 2, 3, 0).clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu()
report_memory("pixels_cpu")
args.output.parent.mkdir(parents=True, exist_ok=True)
raw = args.output.with_suffix(".rgb")
pixels.numpy().tofile(raw)
subprocess.run(["ffmpeg", "-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()
print({"output": str(args.output), "frames": frames, "shape": tuple(pixels.shape)})