import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before any CUDA-touching import
import json
import random
import time
import tempfile

import torch
import gradio as gr
from huggingface_hub import hf_hub_download, snapshot_download

from diffsynth import load_state_dict
from diffsynth.pipelines.wan_video_new import WanVideoPipeline, ModelConfig
from diffsynth.data.video import save_video


# Model repos: T3-Video DiT weights + base Wan2.1-T2V-1.3B components (text encoder, VAE, tokenizer, architecture)
T3_REPO = "APRIL-AIGC/T3-Video"
T3_DIT_FILE = "T3-Video-Wan2.1-T2V-1.3B.safetensors"
BASE_REPO = "Wan-AI/Wan2.1-T2V-1.3B"

# Native 4K output shape (the only resolution the released T3-Video weights support)
HEIGHT, WIDTH, NUM_FRAMES = 2176, 3840, 81
FPS = 30

# Reference negative prompt from the official inference script (infer_multi_gpu.py)
NEGATIVE_PROMPT = (
    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，"
    "低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，"
    "毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
)


def download_weights() -> dict:
    """Fetch all checkpoints from the Hub once at startup.

    Returns a dict of local file paths for the T5 text encoder, the base Wan2.1
    DiT (needed to construct the architecture; its tensors are immediately
    overwritten by the T3-Video state dict), the Wan2.1 VAE and the tokenizer.
    """
    print("[startup] downloading weights from the Hub ...")
    t5 = hf_hub_download(BASE_REPO, "models_t5_umt5-xxl-enc-bf16.pth")
    dit = hf_hub_download(BASE_REPO, "diffusion_pytorch_model.safetensors")
    vae = hf_hub_download(BASE_REPO, "Wan2.1_VAE.pth")
    # the tokenizer needs its config files alongside spiece.model
    tok_dir = snapshot_download(
        BASE_REPO, allow_patterns=["google/umt5-xxl/*"], ignore_patterns=["google/umt5-xxl/*.msgpack", "google/umt5-xxl/*.h5"]
    )
    tok_dir = os.path.join(tok_dir, "google", "umt5-xxl")
    t3_dit = hf_hub_download(T3_REPO, T3_DIT_FILE)
    print("[startup] weights downloaded.")
    return {"t5": t5, "dit": dit, "vae": vae, "tok_dir": tok_dir, "t3_dit": t3_dit}


def load_pipeline() -> WanVideoPipeline:
    """Build the T3-Video pipeline (base Wan2.1-T2V-1.3B components + T3-Video DiT).

    Follows the official infer_multi_gpu.py loading path 1:1: construct from the
    base Wan2.1 files, then overwrite the DiT parameters with the released
    T3-Video full-pretrained state dict. Everything is kept resident on the GPU
    (no CPU offloading; the whole stack fits in 48 GB).
    """
    paths = download_weights()
    model_configs = [
        ModelConfig(path=paths["t5"]),
        ModelConfig(path=paths["dit"]),
        ModelConfig(path=paths["vae"]),
    ]
    pipe = WanVideoPipeline.from_pretrained(
        torch_dtype=torch.bfloat16,
        device="cuda",
        model_configs=model_configs,
        tokenizer_config=ModelConfig(path=paths["tok_dir"]),
        redirect_common_files=False,
    )
    state_dict = load_state_dict(paths["t3_dit"])
    pipe.dit.load_state_dict(state_dict)
    print(f"[startup] loaded T3-Video DiT weights from {paths['t3_dit']}")
    del state_dict
    return pipe


print("[startup] loading pipeline ...")
pipe = load_pipeline()
print("[startup] pipeline ready.")


# --- placeholder constants; retune to the measured per-run time once the live Space is smoke-tested
PER_STEP_SECONDS = 30.0
BASE_SECONDS = 520.0
# PRO visitors get ~2400 s of ZeroGPU quota; requested = duration x 1.5 on this
# hardware, so any single call must stay at or under ~1600 s duration.
MAX_DURATION_SECONDS = 1600


def estimate_duration(prompt: str, steps: int = 15, seed: int = 0, *args, **kwargs) -> int:
    """Estimate the GPU seconds needed for one generation (steps + VAE decode)."""
    # Apply the ~15% headroom BEFORE the cap so large step counts can never
    # request an illegal duration above the visitor's tier cap.
    return int(min(MAX_DURATION_SECONDS, (BASE_SECONDS + PER_STEP_SECONDS * steps) * 1.15))


@spaces.GPU(duration=estimate_duration)
def generate(
    prompt: str,
    steps: int = 15,
    seed: int = 0,
    randomize_seed: bool = True,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a native 4K (2176x3840, 81 frames, 30 fps) video from a text prompt.

    Args:
        prompt: What to generate. Cinematic, camera-motion descriptions work well.
        steps: Denoising steps. 50 matches the paper; 10-20 is much faster.
        seed: RNG seed for reproducibility.
        randomize_seed: Draw a fresh random seed instead of the given one.
    Returns:
        (video file path, used seed, elapsed seconds) for the UI.
    """
    t0 = time.perf_counter()
    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)
    steps = int(steps)

    video = pipe(
        prompt=prompt,
        negative_prompt=NEGATIVE_PROMPT,
        seed=seed,
        height=HEIGHT,
        width=WIDTH,
        num_frames=NUM_FRAMES,
        input_image=None,
        tiled=True,
        tile_size=(68, 120),
        tile_stride=(34, 60),
        cfg_scale=5.0,
        cfg_merge=False,  # separate (not merged) CFG: halves all DiT activations; matches the paper's reference pipeline
        num_inference_steps=steps,
    )
    out_dir = tempfile.mkdtemp()
    out_path = os.path.join(out_dir, "t3_video.mp4")
    save_video(video, out_path, fps=FPS, quality=8)
    elapsed = time.perf_counter() - t0
    print(f"[generate] done in {elapsed:.1f}s (seed={seed}, steps={steps})")
    return out_path, seed, round(elapsed, 1)


# --- curated example prompts from the authors' 4K-World-Vision benchmark (Apache-2.0)
EXAMPLES_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples.json")
with open(EXAMPLES_FILE, "r", encoding="utf-8") as f:
    EXAMPLE_PROMPTS = json.load(f)

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # T3-Video — Native 4K Text-to-Video

            Official demo of **T3-Video-Wan2.1-T2V-1.3B**, a transformer retrofitted with multi-scale
            weight-sharing window attention that generates **native 2176×3840 (4K) video** in one pass —
            no upsampling. [Paper](https://huggingface.co/papers/2512.13492) ·
            [Code](https://github.com/zhangzjn/T3-Video) · [Weights](https://huggingface.co/APRIL-AIGC/T3-Video)

            ⚠️ **Heavy**: every run denoises 685k tokens at 4K. Even at reduced steps one video takes
            **~10–20 minutes of GPU time**, so a Hugging Face PRO account is effectively required
            (free-tier ZeroGPU quota is exhausted by a single run).
            """
        )
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                show_label=False,
                placeholder="Describe the video to generate at native 4K ...",
                container=False,
                scale=4,
            )
            run = gr.Button("Generate", variant="primary", scale=1)
        output = gr.Video(label="Generated 4K video (2176×3840, 81 frames)")
        elapsed = gr.Textbox(label="Elapsed (s)", interactive=False, visible=False)

        with gr.Accordion("Advanced settings", open=False):
            steps = gr.Slider(minimum=10, maximum=50, value=15, step=1, label="Denoising steps",
                              info="50 matches the paper's official setting; 10-20 is much faster.")
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
            seed = gr.Number(label="Seed", value=0, precision=0)

        gr.Examples(
            examples=[[p] for p in EXAMPLE_PROMPTS[:6]],
            inputs=[prompt],
            outputs=[output, seed, elapsed],
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
            label="Example prompts (from the authors' 4K-World-Vision benchmark)",
        )

    run.click(
        generate,
        inputs=[prompt, steps, seed, randomize_seed],
        outputs=[output, seed, elapsed],
        api_name="generate",
    )


if __name__ == "__main__":
    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)