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

def generate_content(api_key, prompt, aspect_ratio, loop):
    """
    Generates video content using LumaAI's API based on user inputs.
    Polls for the video once the generation is completed.

    Parameters:
    - api_key (str): User's LumaAI API key.
    - prompt (str): The prompt for video generation.
    - aspect_ratio (str): Desired aspect ratio (e.g., "16:9").
    - loop (bool): Whether the generation should loop.

    Returns:
    - str: URL to the generated video or an error message.
    """
    try:
        # Initialize the LumaAI client with the provided API key
        client = LumaAI(auth_token=api_key)
        
        # Create a generation request
        generation = client.generations.create(
            aspect_ratio=aspect_ratio,
            loop=loop,
            prompt=prompt,
        )

        # Polling to check if the generation is complete
        generation_id = generation.id
        video_url = None
        
        # Poll until the video is ready
        while video_url is None:
            generation_status = client.generations.get(id=generation_id)
            video_url = generation_status.assets.get('video', None)

        return video_url
    except Exception as e:
        return f"An error occurred: {str(e)}"


def download_video(video_url):
    """
    Downloads the video from the provided URL and saves it locally.

    Parameters:
    - video_url (str): URL of the video to download.

    Returns:
    - str: File path of the downloaded video.
    """
    response = requests.get(video_url, stream=True)
    file_name = 'generated_video.mp4'
    
    with open(file_name, 'wb') as file:
        file.write(response.content)
    
    return file_name

# Create the Gradio Blocks interface
with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# LumaAI Video Generator")

    with gr.Row():
        api_key = gr.Textbox(label="LumaAI API Key", placeholder="Enter your LumaAI API Key", type="password")
    
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", placeholder="Describe what you want to generate...", lines=2)
    
    with gr.Row():
        aspect_ratio = gr.Dropdown(label="Aspect Ratio", choices=["16:9", "4:3", "1:1", "21:9"], value="16:9")
    
    with gr.Row():
        loop = gr.Checkbox(label="Loop", value=False)

    with gr.Row():
        generate_button = gr.Button("Generate Video")
    
    with gr.Row():
        output = gr.Video(label="Generated Video")

    # Define the button click behavior
    def handle_generation(api_key, prompt, aspect_ratio, loop):
        video_url = generate_content(api_key, prompt, aspect_ratio, loop)
        if video_url.startswith("http"):
            file_name = download_video(video_url)
            return file_name
        else:
            return video_url

    # Link the button click to the function
    generate_button.click(
        handle_generation,
        inputs=[api_key, prompt, aspect_ratio, loop],
        outputs=output
    )

# Launch the demo
if __name__ == "__main__":
    demo.launch()