import gradio as gr
from lumaai import LumaAI

def generate_content(api_key, prompt, aspect_ratio, loop):
    """
    Generates content using LumaAI's API based on user inputs.

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

    Returns:
    - str: The ID of the generated content 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,
        )
        
        return f"Generation ID: {generation.id}"
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Create the Gradio Blocks interface
with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# LumaAI Content 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")
    
    with gr.Row():
        output = gr.Textbox(label="Generation Result")
    
    # Link the button click to the function
    generate_button.click(
        generate_content,
        inputs=[api_key, prompt, aspect_ratio, loop],
        outputs=output
    )

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