import gradio as gr
import requests
import random
import time

# FAL endpoint for Wan2.1 T2V
FAL_API_URL = "https://api.fal.ai/fal-ai/wan2.1-t2v-1.3b"

# Function to call fal.ai API
def generate_video(prompt, fps, num_frames, aspect_ratio, guidance_scale, seed, token):
    headers = {
        "Authorization": f"Bearer {token}"
    }

    payload = {
        "prompt": prompt,
        "num_frames": num_frames,
        "fps": fps,
        "aspect_ratio": aspect_ratio,
        "guidance_scale": guidance_scale,
        "negative_prompt": "",
        "seed": int(seed)
    }

    # Call API
    response = requests.post(FAL_API_URL, json=payload, headers=headers)
    response.raise_for_status()
    result = response.json()

    # Poll for video URL
    video_url = None
    for _ in range(20):
        time.sleep(2)
        poll = requests.get(result["urls"]["get"], headers=headers)
        poll.raise_for_status()
        status = poll.json()
        if status["status"] == "COMPLETED":
            video_url = status["output"]["video_url"]
            break
        elif status["status"] == "FAILED":
            raise Exception("Video generation failed.")

    if not video_url:
        raise Exception("Timeout waiting for video.")

    return video_url

# Gradio UI
with gr.Blocks(fill_height=True) as demo:
    with gr.Sidebar():
        gr.Markdown("# Wan2.1 T2V Generator")
        gr.Markdown("Sign in to use this model.")
        login_button = gr.LoginButton()

    gr.Markdown("## 🎬 Generate Video from Text")
    with gr.Row():
        with gr.Column():
            text_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt...")
            fps_slider = gr.Slider(1, 24, value=8, label="FPS")
            frames_slider = gr.Slider(8, 48, value=24, label="Number of Frames")
            aspect_ratio = gr.Dropdown(choices=["square", "landscape", "portrait"], value="square", label="Aspect Ratio")
            guidance_scale = gr.Slider(1, 20, value=15, label="Guidance Scale")
            seed_input = gr.Number(value=random.randint(1, 99999), label="Seed")
            generate_button = gr.Button("Generate Video")
        with gr.Column():
            video_output = gr.Video(label="Generated Video")

    # Wire button to function
    generate_button.click(
        fn=generate_video,
        inputs=[text_input, fps_slider, frames_slider, aspect_ratio, guidance_scale, seed_input, login_button],
        outputs=video_output
    )

demo.launch()