import gradio as gr
import torch
from diffusers import AnimateDiffPipeline, DDIMScheduler
from PIL import Image
import tempfile
import imageio
import time
import os

# Load model
model_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
pipe = AnimateDiffPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Estimate time based on device
def estimate_time(num_frames):
    if torch.cuda.is_available():
        time_per_frame = 2.5  # GPU
    else:
        time_per_frame = 12.0  # CPU

    est_seconds = int(num_frames * time_per_frame)
    return f"Estimated time: ~{est_seconds} seconds ({'GPU' if torch.cuda.is_available() else 'CPU'})"

# Video generation function with simulated progress
def generate_video(image, prompt, num_frames, email):
    status = "Generating..."
    progress = 0

    image = image.convert("RGB").resize((512, 512))

    # Simulated progress before real generation
    for i in range(3):
        time.sleep(0.8)
        progress += 30

    # Generate animation frames
    result = pipe(prompt=prompt, image=image, num_frames=num_frames, guidance_scale=7.5)
    frames = result.frames

    # Save video
    video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    imageio.mimsave(video_path, frames, fps=8)

    # Build message
    final_message = f"✅ Done! Generated {num_frames} frames."
    if email:
        final_message += f"\n📧 Notification would be sent to: {email} (simulated)"

    return video_path, final_message, gr.update(value=100), gr.update(visible=True, value=video_path)

# UI layout
with gr.Blocks() as demo:
    gr.Markdown("# 🌀 Image + Prompt to Video Generator")

    with gr.Row():
        image_input = gr.Image(type="pil", label="🖼️ Upload Image")
        prompt_input = gr.Textbox(label="✏️ Describe Motion (Prompt)")

    with gr.Row():
        num_frames_slider = gr.Slider(8, 32, value=16, step=8, label="🎞️ Number of Frames")
        time_output = gr.Textbox(label="⏱️ Estimated Time", interactive=False)

    email_input = gr.Textbox(label="📧 Optional Email (Notify when done)")
    generate_btn = gr.Button("🎬 Generate Video")

    with gr.Row():
        status_output = gr.Textbox(label="🔄 Status", interactive=False)
        progress_bar = gr.Slider(0, 100, value=0, label="📊 Progress", interactive=False)

    video_output = gr.Video(label="🎥 Output Video")
    download_button = gr.File(label="⬇️ Download Video", visible=False)

    # Events
    num_frames_slider.change(fn=estimate_time, inputs=num_frames_slider, outputs=time_output)

    generate_btn.click(
        fn=generate_video,
        inputs=[image_input, prompt_input, num_frames_slider, email_input],
        outputs=[video_output, status_output, progress_bar, download_button]
    )

    # Optional warning if on CPU
    if not torch.cuda.is_available():
        gr.Markdown("⚠️ **Warning: Running on CPU. Generation will be slow!**")

demo.launch()