"""Prism (FrancisRing/Prism) joint video-audio generation on ZeroGPU (xlarge, 96GB)."""

import os
import sys
import time
import types
import random
import tempfile
import subprocess

import torch

# ---------------------------------------------------------------------------
# ZeroGPU does not support torch.compile -> turn every @torch.compile into a no-op
# (must happen before the model code is imported).
# ---------------------------------------------------------------------------
def _no_compile(fn=None, *args, **kwargs):
    if callable(fn):
        return fn
    return lambda f: f


torch.compile = _no_compile

# The DAC audio VAE only needs `audiotools` for training/IO helpers. Stub it so we
# don't have to install descript-audiotools (it pins protobuf<3.20).
_at = types.ModuleType("audiotools")
_at_ml = types.ModuleType("audiotools.ml")


class _AudioSignal:  # never used in the decode path
    def __init__(self, *a, **k):
        raise RuntimeError("audiotools.AudioSignal is not available in this Space")


_at.AudioSignal = _AudioSignal
class _BaseModel:  # plain mixin; DAC gets nn.Module from diffusers' ModelMixin
    pass


_at_ml.BaseModel = _BaseModel
_at.ml = _at_ml
sys.modules.setdefault("audiotools", _at)
sys.modules.setdefault("audiotools.ml", _at_ml)

import spaces  # noqa: E402  (must be imported before CUDA is touched)
import numpy as np  # noqa: E402
import gradio as gr  # noqa: E402
from PIL import Image  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402
from safetensors import safe_open  # noqa: E402
from accelerate import init_empty_weights  # noqa: E402
from accelerate.utils import set_module_tensor_to_device  # noqa: E402

from hymm.models.modules.mova import MOVABridge  # noqa: E402
from hymm.models.modules.wan_video_dit import WanModel  # noqa: E402
from hymm.models.modules.wan_audio_dit import WanAudioModel  # noqa: E402
from hymm.models.modules.interactionv2 import DualTowerConditionalBridge  # noqa: E402
from hymm.models.modules.dac_vae import DAC  # noqa: E402
from hymm.diffusion.pipelines.mova_pipeline import MOVAPipeline  # noqa: E402
from hymm.diffusion.schedulers.flow_match_pair import FlowMatchPairScheduler  # noqa: E402
from diffusers.models.autoencoders import AutoencoderKLWan  # noqa: E402
from transformers import T5TokenizerFast, UMT5EncoderModel  # noqa: E402

REPO_ID = "FrancisRing/Prism"
VARIANT = os.environ.get("PRISM_VARIANT", "preview_alpha")  # or preview_beta (motion)
BASE = "pretrained_models/MOVA-360p"
DTYPE = torch.bfloat16

NEGATIVE_PROMPT = (
    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，"
    "整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指"
)

# ---------------------------------------------------------------------------
# Download only what inference needs: the Prism transformer checkpoint (it holds
# both video experts, the audio tower and the bridge) + text encoder / VAEs.
# The 57GB of base MOVA DiT weights are fully overridden, so they are skipped.
# ---------------------------------------------------------------------------
t0 = time.time()
root = snapshot_download(
    REPO_ID,
    allow_patterns=[
        f"{VARIANT}/*",
        f"{BASE}/model_index.json",
        f"{BASE}/*/config.json",
        f"{BASE}/text_encoder/*",
        f"{BASE}/tokenizer/*",
        f"{BASE}/scheduler/*",
        f"{BASE}/video_vae/*",
        f"{BASE}/audio_vae/*",
    ],
)
base = os.path.join(root, BASE)
print(f"[prism] download done in {time.time() - t0:.0f}s -> {root}")

# --- transformer: build on meta, then stream the checkpoint straight to cuda ---
t0 = time.time()
with init_empty_weights():
    video_dit = WanModel.from_config(WanModel.load_config(base, subfolder="video_dit"))
    video_dit_2 = WanModel.from_config(WanModel.load_config(base, subfolder="video_dit_2"))
    audio_dit = WanAudioModel.from_config(WanAudioModel.load_config(base, subfolder="audio_dit"))
    dual_bridge = DualTowerConditionalBridge.from_config(
        DualTowerConditionalBridge.load_config(base, subfolder="dual_tower_bridge")
    )
    transformer = MOVABridge(
        video_dit=video_dit,
        video_dit_2=video_dit_2,
        audio_dit=audio_dit,
        dual_tower_bridge=dual_bridge,
        boundary_ratio=0.9,
    )

ckpt_file = os.path.join(root, VARIANT, "diffusion_pytorch_model.safetensors")
expected = set(transformer.state_dict().keys())
loaded = set()
with safe_open(ckpt_file, framework="pt", device="cpu") as f:
    for k in f.keys():
        if k not in expected:
            print(f"[prism] unexpected key {k}")
            continue
        set_module_tensor_to_device(transformer, k, "cuda", value=f.get_tensor(k), dtype=DTYPE)
        loaded.add(k)
missing = expected - loaded
if missing:
    raise RuntimeError(f"[prism] {len(missing)} transformer keys missing, e.g. {sorted(missing)[:5]}")
transformer.to("cuda")
transformer.eval().requires_grad_(False)
print(f"[prism] transformer ready ({len(loaded)} tensors) in {time.time() - t0:.0f}s")

# --- frozen components ---
video_vae = AutoencoderKLWan.from_pretrained(base, subfolder="video_vae", torch_dtype=DTYPE).to("cuda").eval()
audio_vae = DAC.from_pretrained(base, subfolder="audio_vae", torch_dtype=DTYPE).to("cuda").eval()
text_encoder = UMT5EncoderModel.from_pretrained(base, subfolder="text_encoder", torch_dtype=DTYPE).to("cuda").eval()
tokenizer = T5TokenizerFast.from_pretrained(base, subfolder="tokenizer")
scheduler = FlowMatchPairScheduler.from_pretrained(base, subfolder="scheduler")

pipe = MOVAPipeline(
    transformer=transformer,
    video_vae=video_vae,
    audio_vae=audio_vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    scheduler=scheduler,
    boundary_ratio=0.9,
    device=torch.device("cuda"),
)
print("[prism] pipeline ready")

RESOLUTIONS = {
    "480p (848×480)": (480, 848),
    "720p (1280×720)": (720, 1280),
}


def configure_sparse(sparsity, use_ivpq):
    transformer.set_sparse_high_noise_only(False)
    transformer.configure_bsa(
        enable_bsa=True,
        bsa_params={
            "sparsity": float(sparsity),
            "cdf_threshold": 0.20,
            "chunk_3d_shape_q": [4, 4, 4],
            "chunk_3d_shape_k": [4, 4, 4],
        },
        enable_bsa_v2a=False,
        bsa_params_v2a=None,
    )
    for sa in transformer._video_self_attns():
        sa.bsa_params["dynamic_block_lambda_a"] = 0.5
        sa.bsa_params["dynamic_block_tau_128"] = 0.25
        sa.bsa_params["dynamic_block_lambda_128"] = 1.0
    transformer.configure_ivpq_dynamic_block(enable=bool(use_ivpq))


def crop_and_resize(img, height, width):
    w, h = img.size
    target = width / height
    if w / h > target:
        nw = int(h * target)
        img = img.crop(((w - nw) // 2, 0, (w - nw) // 2 + nw, h))
    elif w / h < target:
        nh = int(w / target)
        img = img.crop((0, (h - nh) // 2, w, (h - nh) // 2 + nh))
    return img.resize((width, height), Image.LANCZOS)


def target_size(image, resolution):
    h, w = RESOLUTIONS[resolution]
    if image is not None and image.height > image.width:  # portrait input -> portrait video
        h, w = w, h
    return h, w


def num_tokens(h, w, frames):
    return ((frames - 1) // 4 + 1) * (h // 16) * (w // 16)


# GPU-time model (seconds per transformer forward), calibrated on the Space:
# 480p/49f/30 steps = 255s, 720p/49f/20 steps = 448s (RTX Pro 6000, sparsity 0.75).
SEC_PER_TOKEN = float(os.environ.get("PRISM_SEC_PER_TOKEN", "1.57e-4"))
SEC_PER_TOKEN2 = float(os.environ.get("PRISM_SEC_PER_TOKEN2", "5.1e-9"))
MAX_DURATION = int(os.environ.get("PRISM_MAX_DURATION", "1500"))


def estimate_seconds(image, resolution, frames, steps, cfg, sparsity):
    h, w = target_size(image, resolution)
    tok = num_tokens(h, w, int(frames))
    per_forward = SEC_PER_TOKEN * tok + SEC_PER_TOKEN2 * tok * tok * (1.0 - float(sparsity) + 0.05)
    passes = 1 if float(cfg) == 1.0 else 2
    return int(40 + int(steps) * passes * per_forward)


def get_duration(image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
                 seed, sparsity, use_ivpq, *args, **kwargs):
    return min(MAX_DURATION, int(estimate_seconds(image, resolution, frames, steps, cfg, sparsity) * 1.15))


@spaces.GPU(size="xlarge", duration=get_duration)
def _generate_gpu(image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
                  seed, sparsity, use_ivpq, progress=gr.Progress(track_tqdm=True)):
    t_start = time.time()
    h, w = target_size(image, resolution)
    ref = crop_and_resize(image.convert("RGB"), h, w)
    configure_sparse(sparsity, use_ivpq)
    torch.manual_seed(int(seed))
    video, audio = pipe(
        prompt=prompt,
        image=ref,
        audio_prompt=audio_prompt or prompt,
        negative_prompt=NEGATIVE_PROMPT,
        seed=int(seed),
        height=h,
        width=w,
        num_frames=int(frames),
        video_fps=24.0,
        num_inference_steps=int(steps),
        visual_shift=float(shift),
        audio_shift=float(audio_shift),
        cfg_scale=float(cfg),
        enable_vae_tiling=(h * w >= 720 * 1280),
    )
    frames_np = np.stack([np.asarray(fr) for fr in video[0]])
    audio_np = audio[0].float().cpu().squeeze().numpy()
    elapsed = time.time() - t_start
    print(f"[prism] {h}x{w} {frames}f {steps} steps cfg={cfg} -> {elapsed:.1f}s on GPU "
          f"(peak mem {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB)")
    return frames_np, audio_np, elapsed


def save_video_with_audio(frames, audio_np, path, fps, sample_rate):
    import wave
    import imageio

    with tempfile.TemporaryDirectory() as tmp:
        tmp_video = os.path.join(tmp, "video.mp4")
        tmp_audio = os.path.join(tmp, "audio.wav")
        writer = imageio.get_writer(tmp_video, fps=fps, quality=9, macro_block_size=8)
        for fr in frames:
            writer.append_data(fr)
        writer.close()

        if audio_np.ndim == 1:
            audio_np = audio_np[None, :]
        audio_np = audio_np[:2]
        pcm = (np.clip(audio_np, -1.0, 1.0) * 32767.0).astype(np.int16)
        with wave.open(tmp_audio, "wb") as wf:
            wf.setnchannels(pcm.shape[0])
            wf.setsampwidth(2)
            wf.setframerate(int(sample_rate))
            wf.writeframes(pcm.T.reshape(-1).tobytes())

        cmd = ["ffmpeg", "-y", "-i", tmp_video, "-i", tmp_audio, "-c:v", "copy", "-c:a", "aac",
               "-b:a", "192k", "-movflags", "+faststart", "-shortest", path]
        r = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if r.returncode != 0:
            print("[prism] ffmpeg mux failed:", r.stderr.decode(errors="ignore")[-500:])
            import shutil
            shutil.copyfile(tmp_video, path)


def generate(image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
             seed, randomize_seed, sparsity, use_ivpq, progress=gr.Progress(track_tqdm=True)):
    if image is None:
        raise gr.Error("Please upload a reference image (Prism is an image+text → video+audio model).")
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")
    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
    frames = int(frames)
    frames = frames - (frames - 1) % 4
    est = estimate_seconds(image, resolution, frames, steps, cfg, sparsity)
    if est * 1.15 > MAX_DURATION:
        raise gr.Error(f"This setting needs ~{est / 60:.0f} min of GPU time, above the {MAX_DURATION // 60} min "
                       "per-call limit. Lower frames/steps or use 480p.")
    frames_np, audio_np, elapsed = _generate_gpu(
        image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
        seed, sparsity, use_ivpq,
    )
    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    save_video_with_audio(frames_np, audio_np, out, fps=24.0, sample_rate=pipe.audio_sample_rate)
    return out, seed, f"GPU time: {elapsed:.0f}s (xlarge quota counts ×2)"


def update_estimate(image, resolution, frames, steps, cfg, sparsity):
    h, w = target_size(image, resolution)
    est = estimate_seconds(image, resolution, int(frames), steps, cfg, sparsity)
    return (f"Output {w}×{h}, {int(frames)} frames ({int(frames) / 24:.1f}s) · "
            f"rough GPU estimate ≈ {est / 60:.1f} min → ≈ {2 * est / 60:.1f} min of daily ZeroGPU quota")


EXAMPLES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
examples = []
if os.path.isdir(EXAMPLES_DIR):
    for name in sorted(os.listdir(EXAMPLES_DIR)):
        d = os.path.join(EXAMPLES_DIR, name)
        img, txt = os.path.join(d, "reference_image.jpg"), os.path.join(d, "prompt.txt")
        if os.path.exists(img) and os.path.exists(txt):
            with open(txt, encoding="utf-8") as fh:
                examples.append([img, fh.read().strip()])

with gr.Blocks(title="Prism · joint video + audio") as demo:
    gr.Markdown(
        "# Prism: image + text → video with audio\n"
        f"[Model](https://huggingface.co/{REPO_ID}) · [Code](https://github.com/Tencent-Hunyuan/Prism) · "
        f"[Paper](https://arxiv.org/abs/2610.05416). Checkpoint: `{VARIANT}`. Runs on one ZeroGPU **xlarge** "
        "(RTX Pro 6000, 96GB). xlarge uses quota at 2×, so keep clips short. "
        "Sound can be tagged in the prompt with `<music>`, `<sfx>`, `<speech>`."
    )
    with gr.Row():
        with gr.Column():
            image = gr.Image(label="Reference image (first frame)", type="pil")
            prompt = gr.Textbox(label="Prompt (video)", lines=6)
            audio_prompt = gr.Textbox(label="Audio prompt (optional, defaults to the video prompt)", lines=2)
            with gr.Row():
                resolution = gr.Radio(list(RESOLUTIONS), value="720p (1280×720)", label="Resolution")
                frames = gr.Slider(25, 121, value=49, step=4, label="Frames (24 fps)")
            steps = gr.Slider(10, 50, value=20, step=1, label="Denoising steps")
            estimate = gr.Markdown()
            with gr.Accordion("Advanced", open=False):
                cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="CFG scale (1.0 = no CFG, ~2× faster)")
                shift = gr.Slider(1.0, 17.0, value=7.0, step=0.5, label="Visual shift (720p: 7, 480p: 5)")
                audio_shift = gr.Slider(1.0, 15.0, value=7.0, step=0.5, label="Audio shift")
                sparsity = gr.Dropdown([0.75, 0.85, 0.93], value=0.75, label="Block-sparse attention sparsity")
                use_ivpq = gr.Checkbox(value=True, label="Dynamic block shape (IVPQ)")
                with gr.Row():
                    seed = gr.Number(value=42, precision=0, label="Seed")
                    randomize_seed = gr.Checkbox(value=False, label="Random seed")
            run = gr.Button("Generate", variant="primary")
        with gr.Column():
            out_video = gr.Video(label="Result (with audio)")
            used_seed = gr.Number(label="Seed used", interactive=False)
            info = gr.Markdown()

    est_inputs = [image, resolution, frames, steps, cfg, sparsity]
    for c in est_inputs:
        c.change(update_estimate, est_inputs, estimate)
    demo.load(update_estimate, est_inputs, estimate)

    run.click(
        generate,
        [image, prompt, audio_prompt, resolution, frames, steps, cfg, shift, audio_shift,
         seed, randomize_seed, sparsity, use_ivpq],
        [out_video, used_seed, info],
    )
    if examples:
        gr.Examples(examples, [image, prompt])

demo.queue(max_size=20).launch()