import gradio as gr
import os
from lumaai import LumaAI
import time
import requests

def generate_video(api_key, prompt, loop=False, aspect_ratio="16:9"):
    client = LumaAI(auth_token=api_key)
    
    generation = client.generations.create(
        prompt=prompt,
        loop=loop,
        aspect_ratio=aspect_ratio
    )
    
    # Poll for completion
    while True:
        status = client.generations.get(id=generation.id)
        if status.status == "completed":
            break
        time.sleep(5)
    
    # Download the video
    video_url = status.assets.video
    response = requests.get(video_url, stream=True)
    
    file_name = f"luma_ai_generated_{generation.id}.mp4"
    with open(file_name, 'wb') as file:
        file.write(response.content)
    
    return file_name

def text_to_video(api_key, prompt, loop, aspect_ratio):
    if not api_key:
        return "Please enter your Luma AI API key."
    
    try:
        video_path = generate_video(api_key, prompt, loop, aspect_ratio)
        return video_path
    except Exception as e:
        return f"An error occurred: {str(e)}"

def image_to_video(api_key, prompt, image_url, loop, aspect_ratio):
    if not api_key:
        return "Please enter your Luma AI API key."
    
    try:
        client = LumaAI(auth_token=api_key)
        generation = client.generations.create(
            prompt=prompt,
            loop=loop,
            aspect_ratio=aspect_ratio,
            keyframes={
                "frame0": {
                    "type": "image",
                    "url": image_url
                }
            }
        )
        
        # Poll for completion
        while True:
            status = client.generations.get(id=generation.id)
            if status.status == "completed":
                break
            time.sleep(5)
        
        # Download the video
        video_url = status.assets.video
        response = requests.get(video_url, stream=True)
        
        file_name = f"luma_ai_generated_{generation.id}.mp4"
        with open(file_name, 'wb') as file:
            file.write(response.content)
        
        return file_name
    except Exception as e:
        return f"An error occurred: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# Luma AI Text-to-Video Demo")
    
    api_key = gr.Textbox(label="Luma AI API Key", type="password")
    
    with gr.Tab("Text to Video"):
        prompt = gr.Textbox(label="Prompt")
        generate_btn = gr.Button("Generate Video")
        video_output = gr.Video(label="Generated Video")
        
        with gr.Accordion("Advanced Options", open=False):
            loop = gr.Checkbox(label="Loop", value=False)
            aspect_ratio = gr.Dropdown(label="Aspect Ratio", choices=["16:9", "1:1", "9:16", "4:3", "3:4"], value="16:9")
        
        generate_btn.click(
            text_to_video,
            inputs=[api_key, prompt, loop, aspect_ratio],
            outputs=video_output
        )
    
    with gr.Tab("Image to Video"):
        img_prompt = gr.Textbox(label="Prompt")
        img_url = gr.Textbox(label="Image URL")
        img_generate_btn = gr.Button("Generate Video from Image")
        img_video_output = gr.Video(label="Generated Video")
        
        with gr.Accordion("Advanced Options", open=False):
            img_loop = gr.Checkbox(label="Loop", value=False)
            img_aspect_ratio = gr.Dropdown(label="Aspect Ratio", choices=["16:9", "1:1", "9:16", "4:3", "3:4"], value="16:9")
        
        img_generate_btn.click(
            image_to_video,
            inputs=[api_key, img_prompt, img_url, img_loop, img_aspect_ratio],
            outputs=img_video_output
        )

demo.launch()