"""Run prompt-only H3 T2VA with Ulysses or TP+sequence parallelism."""

import argparse
import json
import os
import time
from pathlib import Path

import torch
import torch.distributed as dist

from h3_blackwell_runtime.checkpoint import H3Checkpoint
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
from h3_blackwell_runtime.distributed import SequenceParallelContext
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.tensor_parallel import configure_h3_tensor_parallel
from h3_blackwell_runtime.token_refiner import H3TokenRefiner


class ParallelDenoiser:
    def __init__(self, model: H3PackedDenoiser, mode: str):
        self.model = model
        self.mode = mode

    def __call__(self, hidden, timesteps, positions, segments, video_segment, audio_segment):
        context = SequenceParallelContext.create(hidden.shape[0], heads=56, head_dim=128)
        if self.mode == "ulysses":
            return self.model.forward_sequence_parallel(
                hidden, timesteps, positions, segments, video_segment, audio_segment, context,
            )
        return self.model.forward_tensor_parallel(
            hidden, timesteps, positions, segments, video_segment, audio_segment, context,
        )


def synchronize(device: torch.device) -> None:
    torch.cuda.synchronize(device)


parser = argparse.ArgumentParser()
parser.add_argument("--benchmark", type=Path, required=True)
parser.add_argument("--mode", choices=("ulysses", "tensor"), required=True)
parser.add_argument("--attention", default="sdpa")
parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
parser.add_argument("--text-encoder", default="/text-encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors")
parser.add_argument("--save-latent", type=Path)
parser.add_argument("--report", type=Path)
args = parser.parse_args()

local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
device = torch.device(f"cuda:{local_rank}")
dist.init_process_group("nccl", device_id=device)
rank = dist.get_rank()
world_size = dist.get_world_size()
benchmark = json.loads(args.benchmark.read_text(encoding="utf-8"))

load_started = time.perf_counter()
checkpoint = H3Checkpoint(args.model, device=device)
model = H3PackedDenoiser.from_checkpoint(
    checkpoint, output_dtype=torch.bfloat16, attention_backend=args.attention,
).eval()
packer = H3PromptPacker(checkpoint)
if args.mode == "tensor":
    partition_context = SequenceParallelContext.create(world_size, heads=56, head_dim=128)
    configure_h3_tensor_parallel(model, partition_context)
synchronize(device)
model_load_seconds = time.perf_counter() - load_started
model_load_peak = torch.cuda.max_memory_allocated(device)
torch.cuda.reset_peak_memory_stats(device)

conditioning_started = time.perf_counter()
if rank == 0:
    conditioner = Qwen3VLPromptConditioner(args.text_encoder, device=device, dtype=torch.float32)
    refiner = H3TokenRefiner(checkpoint, attention_backend="sdpa").eval()
    text = refiner(conditioner(benchmark["prompt"])).to(torch.bfloat16)
    text_length = torch.tensor([text.shape[1]], dtype=torch.int64, device=device)
    del conditioner, refiner
else:
    text = None
    text_length = torch.zeros(1, dtype=torch.int64, device=device)
checkpoint.release_cache()
torch.cuda.empty_cache()
dist.broadcast(text_length, src=0)
if rank != 0:
    text = torch.empty(1, int(text_length.item()), 5376, dtype=torch.bfloat16, device=device)
dist.broadcast(text, src=0)
synchronize(device)
conditioning_seconds = time.perf_counter() - conditioning_started
conditioning_peak = torch.cuda.max_memory_allocated(device)
torch.cuda.reset_peak_memory_stats(device)

width, height = benchmark["resolution"]
video, audio, aligned_frames = random_av_latents(
    width, height, benchmark["frames"], benchmark["seed"], device=device,
)
dist.barrier()
synchronize(device)
sampling_started = time.perf_counter()
video, audio = sample_video_res_multistep(
    ParallelDenoiser(model, args.mode),
    packer,
    text,
    video,
    audio,
    steps=benchmark["steps"],
    seed=benchmark["seed"],
    return_audio=True,
    progress=rank == 0,
)
synchronize(device)
sampling_seconds = time.perf_counter() - sampling_started

timing = torch.tensor(
    [model_load_seconds, conditioning_seconds, sampling_seconds],
    dtype=torch.float64,
    device=device,
)
timings = [torch.empty_like(timing) for _ in range(world_size)]
dist.all_gather(timings, timing)
checksums = torch.stack((video.float().sum(), audio.float().sum())).to(torch.float64)
all_checksums = [torch.empty_like(checksums) for _ in range(world_size)]
dist.all_gather(all_checksums, checksums)
checksum_stack = torch.stack(all_checksums)
if not torch.allclose(checksum_stack, checksum_stack[0].expand_as(checksum_stack), rtol=0, atol=1e-5):
    raise RuntimeError(f"rank outputs diverged: {checksum_stack.cpu().tolist()}")

peak_memory = torch.tensor(
    [model_load_peak, conditioning_peak, torch.cuda.max_memory_allocated(device)],
    dtype=torch.int64,
    device=device,
)
memory = [torch.empty_like(peak_memory) for _ in range(world_size)]
dist.all_gather(memory, peak_memory)

if rank == 0:
    if args.save_latent is not None:
        args.save_latent.parent.mkdir(parents=True, exist_ok=True)
        torch.save({
            "latent": video.cpu(),
            "audio_latent": audio.cpu(),
            "frames": aligned_frames,
            "width": width,
            "height": height,
            "prompt": benchmark["prompt"],
            "seed": benchmark["seed"],
            "distributed_mode": args.mode,
            "world_size": world_size,
            "attention": args.attention,
        }, args.save_latent)
    timing_stack = torch.stack(timings).cpu()
    report = {
        "mode": args.mode,
        "world_size": world_size,
        "attention": args.attention,
        "device": torch.cuda.get_device_name(device),
        "torch": torch.__version__,
        "benchmark": str(args.benchmark),
        "resolution": [width, height],
        "frames": aligned_frames,
        "steps": benchmark["steps"],
        "seed": benchmark["seed"],
        "timings_max_rank_seconds": {
            "model_load": float(timing_stack[:, 0].max()),
            "conditioning": float(timing_stack[:, 1].max()),
            "sampling": float(timing_stack[:, 2].max()),
        },
        "peak_allocated_bytes_by_rank": {
            "model_load": [int(value[0].item()) for value in memory],
            "conditioning": [int(value[1].item()) for value in memory],
            "sampling": [int(value[2].item()) for value in memory],
        },
        "checksums": checksum_stack[0].cpu().tolist(),
        "latent": str(args.save_latent) if args.save_latent is not None else None,
    }
    serialized = json.dumps(report, indent=2)
    if args.report is not None:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(serialized + "\n", encoding="utf-8")
    print(serialized)

dist.destroy_process_group()