import spaces

import os
import uuid
import tempfile

import torch
import gradio as gr

from diffusers import (
    AutoencoderKLWan,
    WanPipeline
)

from diffusers.schedulers.scheduling_unipc_multistep import (
    UniPCMultistepScheduler
)

from diffusers.utils import export_to_video


# ============================================================
# CONFIGURATION
# ============================================================

MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"

OUTPUT_DIR = "/tmp/learnmate_videos"

os.makedirs(OUTPUT_DIR, exist_ok=True)


# ============================================================
# STARTUP INFORMATION
# ============================================================

print("=" * 60)
print("LearnMate Video Generator")
print("=" * 60)

print("Model:", MODEL_ID)
print("PyTorch:", torch.__version__)

print("CUDA available:",
      torch.cuda.is_available())


# ============================================================
# LOAD WAN MODEL
# ============================================================
#
# IMPORTANT:
# ZeroGPU requires the model to be placed on CUDA
# at module level.
#
# Hugging Face ZeroGPU emulates CUDA during startup,
# and provides the real GPU when @spaces.GPU is called.
# ============================================================

print("Loading Wan VAE...")

vae = AutoencoderKLWan.from_pretrained(
    MODEL_ID,
    subfolder="vae",
    torch_dtype=torch.float32
)

print("Loading Wan pipeline...")

pipe = WanPipeline.from_pretrained(
    MODEL_ID,
    vae=vae,
    torch_dtype=torch.bfloat16
)

# 480p recommended flow shift
pipe.scheduler = UniPCMultistepScheduler.from_config(
    pipe.scheduler.config,
    flow_shift=3.0
)

# Move model to CUDA
pipe.to("cuda")

print("Wan model loaded successfully!")

print("=" * 60)


# ============================================================
# VIDEO GENERATION
# ============================================================

@spaces.GPU(duration=180)
def generate_video(prompt):

    if not prompt or not prompt.strip():

        raise gr.Error(
            "Please enter a video description."
        )

    print("=" * 60)
    print("Generating video...")
    print("Prompt:", prompt)

    # --------------------------------------------------------
    # Negative prompt
    # --------------------------------------------------------

    negative_prompt = """
    blurry,
    low quality,
    distorted,
    static image,
    still picture,
    subtitles,
    text,
    watermark,
    logo,
    deformed,
    disfigured,
    ugly,
    extra limbs,
    extra fingers,
    poorly drawn hands,
    poorly drawn faces,
    messy background,
    duplicate objects
    """

    # --------------------------------------------------------
    # Generate video
    # --------------------------------------------------------

    try:

        output = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,

            # 480p
            height=480,
            width=832,

            # Short video
            num_frames=33,

            # Lower number = faster
            num_inference_steps=20,

            guidance_scale=5.0
        )

        frames = output.frames[0]

        # ----------------------------------------------------
        # Save MP4
        # ----------------------------------------------------

        video_id = str(uuid.uuid4())

        video_path = os.path.join(
            OUTPUT_DIR,
            f"{video_id}.mp4"
        )

        export_to_video(
            frames,
            video_path,
            fps=16
        )

        print("Video generated:")
        print(video_path)

        print("=" * 60)

        return video_path

    except Exception as e:

        print("VIDEO GENERATION ERROR:")
        print(str(e))

        raise gr.Error(
            f"Video generation failed: {str(e)}"
        )

# ============================================================
# GRADIO INTERFACE
# ============================================================

description = """
Generate short educational animations from a text description
using Wan2.1 T2V 1.3B.

Designed for the LearnMate personalized learning platform.
"""

examples = [
    [
        "Educational animation showing photosynthesis. "
        "A green plant absorbs sunlight through its leaves, "
        "takes water from the soil through its roots, "
        "and releases oxygen."
    ],

    [
        "Simple educational animation explaining the water cycle. "
        "Water evaporates from the ocean, forms clouds, "
        "then falls as rain and returns to rivers."
    ],

    [
        "Educational animation showing the solar system. "
        "The Earth slowly revolves around the Sun "
        "while the planets move in their orbits."
    ]
]


demo = gr.Interface(

    fn=generate_video,

    inputs=gr.Textbox(
        label="Learning Video Prompt",
        placeholder="Describe the educational animation...",
        lines=5
    ),

    outputs=gr.Video(
        label="Generated Learning Video",
        format="mp4",
        autoplay=False
    ),

    title="LearnMate 🎓 → 🎬",

    description=description,

    examples=examples,

    api_name="generate_video"
)


# ============================================================
# LAUNCH
# ============================================================

demo.launch()