import os
import tempfile

import spaces  # must come before torch / any CUDA-touching import

import numpy as np
import torch
import gradio as gr
import decord
import imageio
from PIL import Image
from omegaconf import OmegaConf
from torchvision import transforms
from einops import rearrange
from huggingface_hub import snapshot_download, hf_hub_download

from pipeline import CausalInferencePipeline
from utils.misc import set_seed

# ----------------------------------------------------------------------------
# Constants (mirror the released `infer-local-ar-forcing.sh` defaults)
# ----------------------------------------------------------------------------
WAN_REPO = "Wan-AI/Wan2.1-T2V-1.3B"
WAN_DIR = "wan_models/Wan2.1-T2V-1.3B"
LIVEEDIT_REPO = "cp-cp/LiveEdit"
CKPT_NAME = "ar-forcing_002000.pt"
CONFIG_PATH = "configs/wan_mm-ar-forcing-local.yaml"
DEFAULT_CONFIG_PATH = "configs/default_config.yaml"

NUM_OUTPUT_FRAMES = 21          # latent frames -> (21-1)*4 + 1 = 81 pixel frames
NUM_PIXEL_FRAMES = NUM_OUTPUT_FRAMES * 4 - 3
HEIGHT, WIDTH = 480, 832
FPS = 16

# ----------------------------------------------------------------------------
# Download weights (base Wan2.1-T2V-1.3B + the LiveEdit causal checkpoint)
# ----------------------------------------------------------------------------
os.makedirs("wan_models", exist_ok=True)
os.makedirs("checkpoints/liveedit", exist_ok=True)

snapshot_download(
    repo_id=WAN_REPO,
    local_dir=WAN_DIR,
    allow_patterns=[
        "config.json",
        "diffusion_pytorch_model.safetensors",
        "Wan2.1_VAE.pth",
        "models_t5_umt5-xxl-enc-bf16.pth",
        "google/umt5-xxl/*",
    ],
)

CKPT_PATH = hf_hub_download(
    repo_id=LIVEEDIT_REPO,
    filename=CKPT_NAME,
    local_dir="checkpoints/liveedit",
)

# ----------------------------------------------------------------------------
# Build the pipeline at module scope (ZeroGPU streams the eager .cuda weights
# into VRAM on the first @spaces.GPU call).
# ----------------------------------------------------------------------------
device = torch.device("cuda")
set_seed(0)
torch.set_grad_enabled(False)

config = OmegaConf.merge(
    OmegaConf.load(DEFAULT_CONFIG_PATH),
    OmegaConf.load(CONFIG_PATH),
)

# Few-step causal editing pipeline with the source-video conditioning channels
# (patch embedding expanded 16 -> 32) exactly as in inference-mm.py.
local_attn_size = -1  # NUM_OUTPUT_FRAMES (21) is not > 21, so global attention
pipeline = CausalInferencePipeline(
    config,
    device=device,
    local_attn_size=local_attn_size,
    sink_size=0,
    expand_patch_embedding=True,
)


def _remove_fsdp_wrapped_module(state_dict):
    cleaned = {}
    for key, value in state_dict.items():
        if "_fsdp_wrapped_module" in key:
            new_key = "model." + key.split("._fsdp_wrapped_module.")[-1]
            cleaned[new_key] = value
        else:
            cleaned[key] = value
    return cleaned


def _select_generator_state_dict(state_dict):
    if "generator" in state_dict:
        return state_dict["generator"]
    if "generator_ema" in state_dict:
        return state_dict["generator_ema"]
    return state_dict


_raw = torch.load(CKPT_PATH, map_location="cpu")
_gen = _select_generator_state_dict(_raw)
pipeline.generator.load_state_dict(_remove_fsdp_wrapped_module(_gen))

pipeline = pipeline.to(dtype=torch.bfloat16)
pipeline.text_encoder.to(device=device)
pipeline.generator.to(device=device)
pipeline.vae.to(device=device)

_TRANSFORM = transforms.Compose([
    transforms.Resize((HEIGHT, WIDTH)),
    transforms.ToTensor(),
])


def _read_source_video(video_path):
    """Load and resample a source video to NUM_PIXEL_FRAMES at 480x832.

    Returns a tensor of shape [1, C, T, H, W] in [0, 1].
    """
    vr = decord.VideoReader(video_path, ctx=decord.cpu(0))
    total = len(vr)
    if total >= NUM_PIXEL_FRAMES:
        indices = np.arange(0, NUM_PIXEL_FRAMES)
    else:
        indices = np.linspace(0, total - 1, NUM_PIXEL_FRAMES, dtype=int)
    frames = vr.get_batch(indices).asnumpy()  # (T, H, W, C)
    frames = [Image.fromarray(f) for f in frames]
    video = torch.stack([_TRANSFORM(f) for f in frames])  # [T, C, H, W]
    video = video.permute(1, 0, 2, 3).unsqueeze(0)        # [1, C, T, H, W]
    return video


@spaces.GPU(duration=120)
@torch.inference_mode()
def edit_video(video_path, instruction, seed, progress=gr.Progress(track_tqdm=True)):
    if not video_path:
        raise gr.Error("Please provide a source video.")
    if not instruction or not instruction.strip():
        raise gr.Error("Please provide an editing instruction.")

    set_seed(int(seed))

    source_pixel = _read_source_video(video_path).to(device=device, dtype=torch.bfloat16)
    source_latent = pipeline.vae.encode_to_latent(source_pixel).to(dtype=torch.bfloat16)

    noise = torch.randn(
        [1, NUM_OUTPUT_FRAMES, 16, 60, 104], device=device, dtype=torch.bfloat16,
        generator=torch.Generator(device=device).manual_seed(int(seed)),
    )

    video, _ = pipeline.inference(
        noise=noise,
        text_prompts=[instruction.strip()],
        return_latents=True,
        initial_latent=None,
        y=source_latent,
        wo_scale=True,
    )

    pipeline.vae.model.clear_cache()

    # [1, T, C, H, W] in [0, 1] -> uint8 [T, H, W, C]
    frames = rearrange(video, "b t c h w -> b t h w c")[0]
    frames = (frames.float().clamp(0, 1) * 255.0).byte().cpu().numpy()

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    imageio.mimwrite(out_path, list(frames), fps=FPS, codec="libx264", quality=8)
    return out_path


DESCRIPTION = """
# LiveEdit — Real-Time Diffusion-Based Streaming Video Editing

Causal, chunk-by-chunk video editing built on **Wan2.1-T2V-1.3B** and the
Self-Forcing codebase. Give it a source video and a text instruction; LiveEdit
edits the video while preserving backgrounds and non-edited regions.

[Project page](https://live-edit.github.io) ·
[Paper](https://arxiv.org/abs/2606.26740) ·
[Code](https://github.com/cp-cp/LiveEdit) ·
[Checkpoints](https://huggingface.co/cp-cp/LiveEdit) · ECCV 2026

The source video is resampled to 81 frames at 480×832; output is ~5s @ 16fps.
"""

with gr.Blocks(title="LiveEdit") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            in_video = gr.Video(label="Source video", sources=["upload"])
            instruction = gr.Textbox(
                label="Editing instruction",
                placeholder="e.g. Change the red currants to deep black grapes.",
            )
            seed = gr.Slider(0, 2**31 - 1, value=0, step=1, label="Seed")
            run = gr.Button("Edit video", variant="primary")
        with gr.Column():
            out_video = gr.Video(label="Edited video")

    gr.Examples(
        examples=[["test_cases/test.mp4", "Change the red currants to deep black grapes.", 0]],
        inputs=[in_video, instruction, seed],
        outputs=out_video,
        fn=edit_video,
        cache_examples=False,
    )

    run.click(edit_video, inputs=[in_video, instruction, seed], outputs=out_video)

if __name__ == "__main__":
    demo.launch()