import gradio as gr
from lumaai.client import LumaAI  # Adjust based on your findings
import requests
from PIL import Image
from io import BytesIO
import os

# Initialize LumaAI client
def create_client(api_key):
    return LumaAI(auth_token=api_key)

# Basic video generation
def generate_video(api_key, prompt, aspect_ratio, loop):
    client = create_client(api_key)
    try:
        generation = client.generations.create(
            prompt=prompt,
            aspect_ratio=aspect_ratio,
            loop=loop
        )
        return generation.id, "Generation started..."
    except Exception as e:
        return None, f"Error: {str(e)}"

# Polling generation status
def poll_generation(api_key, generation_id):
    if not generation_id:
        return "No generation in progress", None
    
    client = create_client(api_key)
    try:
        generation = client.generations.get(id=generation_id)
        status = generation.status
        thumbnail_url = generation.assets.thumbnail
        
        # Fetch the thumbnail image
        thumbnail_response = requests.get(thumbnail_url)
        thumbnail = thumbnail_response.content if thumbnail_response.status_code == 200 else None
        return status, thumbnail
    except Exception as e:
        return f"Error: {str(e)}", None

# Download the generated video
def download_video(api_key, generation_id):
    if not generation_id:
        return None
    
    client = create_client(api_key)
    try:
        generation = client.generations.get(id=generation_id)
        video_url = generation.assets.video
        response = requests.get(video_url, stream=True)
        
        if response.status_code != 200:
            return None
        
        file_name = f"generated_video_{generation_id}.mp4"
        with open(file_name, 'wb') as file:
            for chunk in response.iter_content(chunk_size=8192):
                if chunk:
                    file.write(chunk)
        
        return file_name
    except Exception as e:
        return None

with gr.Blocks() as demo:
    gr.Markdown("# Luma AI Video Generation Demo")
    gr.Markdown("Generate videos using Luma AI based on text prompts.")
    
    api_key = gr.Textbox(
        label="Enter your Luma AI API Key",
        type="password",
        placeholder="sk-..."
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe your video...",
                lines=2
            )
            aspect_ratio = gr.Dropdown(
                choices=["16:9", "9:16", "1:1", "4:3", "3:4"],
                label="Aspect Ratio",
                value="16:9"
            )
            loop = gr.Checkbox(
                label="Loop",
                value=False
            )
            generate_btn = gr.Button("Generate Video")
        
        with gr.Column(scale=3):
            status = gr.Textbox(
                label="Status",
                interactive=False  # Replaces readonly=True
            )
            thumbnail = gr.Image(
                label="Thumbnail",
                interactive=False
            )
            video = gr.Video(
                label="Generated Video",
                interactive=False
            )
    
    generation_id = gr.State()
    
    # Generate Video Action
    def on_generate(api_key, prompt, aspect_ratio, loop):
        if not api_key:
            return None, "Please provide a valid API key."
        if not prompt:
            return None, "Prompt cannot be empty."
        
        gen_id, message = generate_video(api_key, prompt, aspect_ratio, loop)
        return gen_id, message
    
    generate_btn.click(
        on_generate,
        inputs=[api_key, prompt, aspect_ratio, loop],
        outputs=[generation_id, status]
    )
    
    # Check Status and Display Results
    def on_poll(api_key, gen_id):
        if not api_key:
            return "Please provide a valid API key.", None, None
        if not gen_id:
            return "No generation in progress.", None, None
        
        status_text, thumb_data = poll_generation(api_key, gen_id)
        if status_text.lower() == "completed":
            video_path = download_video(api_key, gen_id)
            if video_path:
                video_output = video_path
            else:
                video_output = None
        else:
            video_output = None
        
        # Convert thumbnail bytes to image format
        if thumb_data:
            thumb_image = Image.open(BytesIO(thumb_data))
        else:
            thumb_image = None
        
        return status_text, thumb_image, video_output
    
    poll_btn = gr.Button("Check Status")
    poll_btn.click(
        on_poll,
        inputs=[api_key, generation_id],
        outputs=[status, thumbnail, video]
    )
    
    gr.Markdown("""
    ---
    ### Documentation & Resources
    - [Luma AI Python SDK](https://github.com/lumalabs/lumaai-python)
    - [Get API Key](https://lumalabs.ai/dream-machine/api/keys)
    """)

demo.launch()