"""
🎬 Image-to-Video AI Studio — Real AI Animation
Uses CogVideoX-5B-I2V diffusion model to generate REAL animated videos from images.
Runs on HuggingFace ZeroGPU (free) — actual AI video generation, not zoom/warp tricks.
"""

import spaces
import torch
import gradio as gr
import tempfile
import gc
from PIL import Image
from diffusers import CogVideoXImageToVideoPipeline
from diffusers.utils import export_to_video, load_image

# ─── Model Loading (CPU at startup, GPU inside @spaces.GPU) ─────────────────
print("📦 Loading CogVideoX-5B-I2V pipeline to CPU...")
pipe = CogVideoXImageToVideoPipeline.from_pretrained(
    "THUDM/CogVideoX-5b-I2V",
    torch_dtype=torch.bfloat16,
)
print("✅ Pipeline loaded to CPU successfully!")


# ─── Prompt Templates ────────────────────────────────────────────────────────
CAMERA_PROMPTS = {
    "None (Auto)": "",
    "Static Shot": "static camera, no camera movement, steady shot",
    "Slow Zoom In": "camera slowly zooms in, push in shot",
    "Slow Zoom Out": "camera slowly zooms out, pull back reveal",
    "Pan Left": "camera pans slowly to the left, horizontal tracking",
    "Pan Right": "camera pans slowly to the right, horizontal tracking",
    "Tilt Up": "camera tilts upward slowly, revealing the sky",
    "Tilt Down": "camera tilts downward slowly",
    "Orbit Shot": "camera orbits around the subject slowly",
    "Dolly Forward": "camera moves forward smoothly, dolly shot",
    "Crane Up": "camera rises upward, crane shot, aerial movement",
    "Handheld": "slight handheld camera movement, natural sway",
}

STYLE_PROMPTS = {
    "Cinematic": "cinematic lighting, film grain, dramatic, movie quality, 4K",
    "Realistic": "photorealistic, natural lighting, high detail, sharp focus",
    "Dreamy": "soft focus, ethereal glow, dreamy atmosphere, gentle lighting",
    "Dynamic": "fast motion, energetic, vibrant, dynamic movement",
    "Slow Motion": "slow motion, graceful movement, smooth, elegant",
    "Nature Documentary": "nature documentary style, wildlife, natural beauty, BBC quality",
    "Anime / Artistic": "anime style, artistic, vibrant colors, stylized movement",
}

DURATION_FRAMES = {
    "~3 seconds (24 frames)": 24,
    "~6 seconds (49 frames)": 49,
}


# ─── Core Video Generation ───────────────────────────────────────────────────
@spaces.GPU(duration=300)
def generate_video(
    image,
    prompt,
    camera_motion,
    style,
    duration,
    num_inference_steps,
    guidance_scale,
    seed,
    use_random_seed,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate real animated video using CogVideoX-5B-I2V diffusion model."""

    if image is None:
        raise gr.Error("⚠️ Please upload an image first!")

    if not prompt or prompt.strip() == "":
        raise gr.Error("⚠️ Please describe what animation you want!")

    # Move pipeline to GPU
    pipe.to("cuda")
    pipe.vae.enable_tiling()
    pipe.vae.enable_slicing()

    # Build the full prompt
    parts = [prompt.strip()]
    camera_desc = CAMERA_PROMPTS.get(camera_motion, "")
    if camera_desc:
        parts.append(camera_desc)
    style_desc = STYLE_PROMPTS.get(style, "")
    if style_desc:
        parts.append(style_desc)
    parts.append("high quality, detailed, smooth motion")
    full_prompt = ". ".join(parts)

    print(f"🎬 Prompt: {full_prompt}")

    # Seed
    if use_random_seed:
        seed = torch.randint(0, 2**32, (1,)).item()
    generator = torch.Generator(device="cuda").manual_seed(int(seed))

    # Frames
    num_frames = DURATION_FRAMES.get(duration, 49)

    # Resize image (CogVideoX native: 720x480 landscape, 480x720 portrait)
    if isinstance(image, str):
        image = load_image(image)
    img_w, img_h = image.size
    if img_w >= img_h:
        target_w, target_h = 720, 480
    else:
        target_w, target_h = 480, 720
    image_resized = image.resize((target_w, target_h), Image.LANCZOS)

    # Generate
    print(f"🚀 Generating: {num_frames} frames, {num_inference_steps} steps")
    result = pipe(
        prompt=full_prompt,
        image=image_resized,
        num_videos_per_prompt=1,
        num_inference_steps=int(num_inference_steps),
        num_frames=num_frames,
        guidance_scale=float(guidance_scale),
        generator=generator,
    ).frames[0]

    # Export to MP4
    out_path = tempfile.mktemp(suffix=".mp4")
    export_to_video(result, out_path, fps=8)

    # Cleanup
    pipe.to("cpu")
    gc.collect()
    torch.cuda.empty_cache()

    print(f"✅ Video saved: {out_path}")
    return out_path, f"🎲 Seed: {seed} | Frames: {num_frames} | Steps: {num_inference_steps}"


# ─── Gradio UI ───────────────────────────────────────────────────────────────
CSS = """
.main-header { text-align: center; padding: 20px 0 10px 0; }
.main-header h1 {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text; -webkit-text-fill-color: transparent;
    font-size: 2.2em; margin-bottom: 5px;
}
.note-box {
    background: linear-gradient(135deg, #fff3cd 0%, #ffeeba 100%);
    border-left: 4px solid #ffc107; border-radius: 8px;
    padding: 12px 16px; margin: 10px 0;
}
.info-box {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    border-radius: 10px; padding: 15px; margin: 10px 0;
}
footer { display: none !important; }
"""

with gr.Blocks(title="🎬 Image-to-Video AI", css=CSS) as demo:

    gr.HTML("""
    <div class="main-header">
        <h1>🎬 Image-to-Video AI</h1>
        <p style="color: #666; font-size: 1.15em;">
            Transform any image into a <b>real animated video</b> with AI
        </p>
        <p style="color: #999; font-size: 0.85em;">
            CogVideoX-5B · Free ZeroGPU · Real Diffusion Animation
        </p>
    </div>
    """)

    gr.HTML("""
    <div class="note-box">
        ⚡ <b>Real AI Animation:</b> This uses a 5-billion parameter diffusion model to generate 
        <b>actual new video frames</b> — not zoom/pan tricks. Water flows, hair blows, 
        people walk, clouds drift. Takes ~2-4 min per video on free GPU.
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(label="📸 Upload Image", type="pil", height=320)

            prompt_input = gr.Textbox(
                label="✍️ Animation Prompt — describe what should move",
                placeholder="e.g., 'The river flows gently, trees sway in the breeze, clouds drift across the sky'",
                lines=3,
            )

            with gr.Row():
                camera_motion = gr.Dropdown(
                    choices=list(CAMERA_PROMPTS.keys()),
                    value="None (Auto)", label="🎥 Camera Motion",
                )
                style = gr.Dropdown(
                    choices=list(STYLE_PROMPTS.keys()),
                    value="Cinematic", label="🎭 Style",
                )

            duration = gr.Radio(
                choices=list(DURATION_FRAMES.keys()),
                value="~6 seconds (49 frames)", label="⏱️ Duration",
            )

            with gr.Accordion("⚙️ Advanced", open=False):
                num_inference_steps = gr.Slider(
                    minimum=10, maximum=50, value=30, step=5,
                    label="🔄 Steps (15=fast, 30=good, 50=best)",
                )
                guidance_scale = gr.Slider(
                    minimum=1.0, maximum=12.0, value=6.0, step=0.5,
                    label="🎯 Guidance Scale",
                )
                with gr.Row():
                    seed = gr.Number(label="🎲 Seed", value=42, precision=0)
                    use_random_seed = gr.Checkbox(label="Random", value=True)

            generate_btn = gr.Button("🎬 Generate Video", variant="primary", size="lg")

        with gr.Column(scale=1):
            video_output = gr.Video(label="🎥 Generated Video", autoplay=True, loop=True, height=420)
            gen_info = gr.Textbox(label="Info", interactive=False)

            gr.HTML("""
            <div class="info-box">
                <b>🧠 How it works:</b>
                <ul style="margin: 5px 0; padding-left: 18px; color: #555;">
                    <li>CogVideoX-5B generates <b>real new video frames</b> using diffusion</li>
                    <li>Objects actually move — water, people, animals, particles</li>
                    <li>Your prompt controls <b>what</b> animates and <b>how</b></li>
                    <li>Camera motion and style are added to the AI prompt</li>
                </ul>
                <p style="color: #888; font-size: 0.85em; margin-top: 8px;">
                    <a href="https://huggingface.co/THUDM/CogVideoX-5b-I2V" target="_blank">CogVideoX-5B-I2V</a> 
                    by Tsinghua University · ZeroGPU (Free A10G)
                </p>
            </div>
            """)

    generate_btn.click(
        fn=generate_video,
        inputs=[image_input, prompt_input, camera_motion, style, duration,
                num_inference_steps, guidance_scale, seed, use_random_seed],
        outputs=[video_output, gen_info],
    )

    gr.Markdown("### 💡 Example Prompts")
    gr.Markdown("""
| Image Type | Example Prompt |
|-----------|---------------|
| 🏞️ Landscape | *"The river flows gently, clouds drift, trees sway in breeze"* |
| 👤 Portrait | *"She smiles and turns her head, hair flowing in the wind"* |
| 🌊 Ocean | *"Waves crash on shore, sea foam bubbles, seagulls fly"* |
| 🏙️ City | *"Cars drive, pedestrians walk, city lights flicker"* |
| 🐾 Animals | *"The cat stretches, yawns, then walks forward"* |
| 🌸 Nature | *"Petals sway in breeze, butterfly lands on flower"* |
    """)

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