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

import spaces  # MUST come before any torch / CUDA-touching import
import torch
import gradio as gr
import imageio
import json
import tempfile
import time
from pathlib import Path

import safetensors.torch
from huggingface_hub import snapshot_download
from diffusers import FlowMatchEulerDiscreteScheduler, WanPipeline

from mobilewan.config import (
    REHYAT_NUM_BLOCKS,
    SAMPLE_CFG_SCALE,
    SAMPLE_HEIGHT,
    SAMPLE_NUM_FRAMES,
    SAMPLE_WIDTH,
)
from mobilewan.rehyat_utils import (
    apply_pruning_plan,
    load_pruning_plan,
    surgery,
)

BASE_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
MOBILEWAN_CHECKPOINT_REPO = "Qualcomm-AI-Research/mobilewan"
PRUNING_PLAN_PATH = Path(__file__).parent / "mobilewan" / "pruning_plan.json"
VIDEO_FPS = 16


class FlowMatchEulerDiscreteLowStepScheduler(FlowMatchEulerDiscreteScheduler):
    """FlowMatch Euler scheduler variant that drops the final timestep for low-step sampling."""

    def set_timesteps(
        self,
        num_inference_steps = None,
        device = None,
    ):
        super().set_timesteps(num_inference_steps=num_inference_steps + 1, device=device)
        self.timesteps = self.timesteps[:-1]  # remove last timestep == 1
        self.sigmas = torch.cat([self.sigmas[:-2], self.sigmas[-1:]])


def _load_transformer_assign(model, ckpt_dir):
    """Load checkpoint weights using assign=True to avoid inference-mode tensor errors.

    The pruning helpers run under @torch.inference_mode(), which creates
    inference-mode parameters. In-place copies (as done by safetensors.load_model
    / model.load_state_dict) fail on those tensors. Using assign=True replaces
    the parameters entirely rather than copying in-place, sidestepping the issue.
    """
    ckpt_dir = Path(ckpt_dir)
    single_file = ckpt_dir / "diffusion_pytorch_model.safetensors"
    index_file = ckpt_dir / "diffusion_pytorch_model.safetensors.index.json"

    if single_file.exists():
        state_dict = safetensors.torch.load_file(str(single_file))
    elif index_file.exists():
        with index_file.open("r") as f:
            index = json.load(f)
        state_dict = {}
        for shard_file in sorted(set(index["weight_map"].values())):
            shard_path = ckpt_dir / shard_file
            state_dict.update(safetensors.torch.load_file(str(shard_path)))
    else:
        raise FileNotFoundError(f"Could not find model weights in {ckpt_dir}")

    model.load_state_dict(state_dict, assign=True)


print("Downloading MobileWan checkpoint...")
mobilewan_ckpt_dir = snapshot_download(
    MOBILEWAN_CHECKPOINT_REPO, repo_type="model"
)

print("Loading base WanPipeline...")
pipe = WanPipeline.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.bfloat16)

start_block = min(2, 30 - REHYAT_NUM_BLOCKS)
block_inds = list(range(start_block, REHYAT_NUM_BLOCKS + start_block))
surgery(pipe.transformer, block_inds)

pruning_plan = apply_pruning_plan(
    pipe.transformer,
    load_pruning_plan(PRUNING_PLAN_PATH),
)
print(f"Applied pruning plan: {pruning_plan}")

print("Loading MobileWan checkpoint...")
_load_transformer_assign(pipe.transformer, mobilewan_ckpt_dir)

pipe.vae.encoder = torch.nn.Identity()
pipe.scheduler = FlowMatchEulerDiscreteLowStepScheduler(shift=5)
pipe.set_progress_bar_config(disable=False)
pipe.to("cuda")
pipe.transformer.to(torch.bfloat16)
print("Model loaded and ready.")


def _estimate(prompt, num_steps, *args, **kwargs):
    # Swallow extras with *args, **kwargs — Gradio passes progress= positionally
    return min(240, 25 + int(num_steps) * 8)


@spaces.GPU(duration=_estimate)
def generate(
    prompt: str,
    seed: int = 0,
    num_steps: int = 3,
    fps: int = VIDEO_FPS,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a short video from a text prompt using MobileWan.

    Args:
        prompt: Text prompt describing the video to generate.
        seed: RNG seed for reproducible sampling. 0 means fixed seed.
        num_steps: Number of denoising steps (default 3 for the distilled model).
        fps: Output video frame rate.
    """
    seed = int(seed)
    if seed == 0:
        seed = 42
    generator = torch.Generator(device="cuda").manual_seed(seed)

    t0 = time.perf_counter()
    videos = pipe(
        [prompt],
        num_inference_steps=int(num_steps),
        guidance_scale=SAMPLE_CFG_SCALE,
        num_frames=SAMPLE_NUM_FRAMES,
        height=SAMPLE_HEIGHT,
        width=SAMPLE_WIDTH,
        generator=generator,
        output_type="pt",
    ).frames
    elapsed = time.perf_counter() - t0

    video = videos[0]
    vid = torch.einsum("tchw->thwc", video.mul(255).clamp(0, 255).to(torch.uint8)).cpu().numpy()

    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    writer = imageio.get_writer(out_path, fps=int(fps), codec="libx264", quality=5)
    try:
        for frame in vid:
            writer.append_data(frame)
    finally:
        writer.close()

    return out_path, f"Generated in {elapsed:.1f}s (seed={seed}, {num_steps} steps)"


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

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # MobileWan: Closing the Quality Gap for Mobile Video Diffusion
        Generate a short video (~5s, 480×832) from a text prompt using Qualcomm AI Research's MobileWan — a lightweight, mobile-optimized adaptation of Wan2.2-5B with hybrid attention and head pruning.
        """
    )

    with gr.Row():
        prompt = gr.Textbox(
            label="Prompt",
            show_label=False,
            placeholder="Describe the video you want to generate...",
            container=False,
            scale=4,
        )
        run = gr.Button("Generate", variant="primary", scale=1)

    video_out = gr.Video(label="Generated Video")
    status = gr.Textbox(label="Status", interactive=False)

    with gr.Accordion("Advanced Settings", open=False):
        seed = gr.Number(label="Seed", value=0, precision=0, info="0 = fixed default seed")
        num_steps = gr.Slider(
            label="Denoising Steps",
            minimum=1,
            maximum=10,
            step=1,
            value=3,
            info="3 is the default for the distilled MobileWan model",
        )
        fps = gr.Slider(
            label="Output FPS",
            minimum=8,
            maximum=30,
            step=1,
            value=VIDEO_FPS,
        )

    gr.Examples(
        examples=[
            ["A cinematic shot of a golden retriever running through a sunny backyard, weaving between flower beds and patches of tall grass. Warm afternoon light filters through the trees, creating soft shadows and a playful, energetic atmosphere."],
            ["A battle-hardened warrior, face and body smeared with mud and fresh wounds, wearing heavy leather armor and furs. In a heavy downpour on a muddy battlefield, he takes a slow step forward and tightens his grip on his weapon, rain streaming off his brow as he glares angrily ahead."],
            ["A serene timelapse of a mountain lake at dawn, mist rolling over the water as the sun rises behind snow-capped peaks, reflecting vibrant orange and pink hues across the calm surface."],
            ["A busy city street at night, neon signs reflecting off wet pavement, pedestrians with umbrellas hurrying past, cars leaving light trails as they drive through the rain."],
        ],
        inputs=[prompt],
        outputs=[video_out, status],
        fn=generate,
        cache_examples=True,
        cache_mode="lazy",
    )

    run.click(
        fn=generate,
        inputs=[prompt, seed, num_steps, fps],
        outputs=[video_out, status],
    )

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