import gradio as gr
import time
import random
import requests
from lumaai import LumaAI
import traceback

from lib.status_utils import StatusTracker, load_messages
from lib.image_utils import prepare_image
from lib.api_utils import get_camera_motions
from lib.ui_components import create_input_column, create_output_column

def generate_video(api_key, prompt, camera_motion, image=None, progress=gr.Progress()):
    status_box = gr.Markdown()  # Create status box
    if not api_key or not prompt:
        return None, "Please provide both API key and prompt (I'm not a mind reader... yet)"
    
    try:
        status_tracker = StatusTracker(progress, status_box)
        status_tracker.add_step("LumaAI initialized", 0.01)
        client = LumaAI(auth_token=api_key)
        
        # Prepare generation parameters
        generation_params = {
            "prompt": f"{prompt} {camera_motion if camera_motion != 'None' else ''}",
            "loop": True,
            "aspect_ratio": "1:1"  # Force square aspect ratio
        }

        # Handle image if provided
        if image is not None:
            try:
                cdn_url = prepare_image(image, status_tracker)
                generation_params["keyframes"] = {
                    "frame0": {
                        "type": "image",
                        "url": cdn_url
                    }
                }
                status_tracker.add_step("Image ready for its starring role", 0.1)
            except Exception as e:
                return None, f"🎭 Drama in the image department: {str(e)}"

        status_tracker.add_step("Sending your creative masterpiece to LumaAI", 0.15)
        try:
            generation = client.generations.create(**generation_params)
        except Exception as e:
            return None, f"🎬 LumaAI didn't like that: {str(e)}"
        
        # Load and shuffle status messages
        status_messages = load_messages()
        random.shuffle(status_messages)
        
        # Poll for completion
        start_time = time.time()
        message_index = 0
        last_status = None
        
        while True:
            try:
                generation_status = client.generations.get(generation.id)
                status = generation_status.state
                elapsed_time = time.time() - start_time
                
                if status != last_status:
                    status_tracker.add_step(f"Status: {status}", min(0.2 + (elapsed_time/300), 0.8))
                    last_status = status
                
                current_message = status_messages[message_index % len(status_messages)]
                status_tracker.update_message(current_message, min(0.2 + (elapsed_time/300), 0.8))
                message_index += 1
                
                if status == 'completed':
                    status_tracker.add_step("Generation completed!", 0.9)
                    download_url = generation_status.assets.video
                    break
                elif status == 'failed':
                    failure_reason = generation_status.failure_reason or "It's not you, it's me"
                    return None, f"🎭 Generation failed: {failure_reason}"
                
                if elapsed_time > 300:
                    return None, "⏰ Generation timeout (5 minutes of awkward silence)"
                
                time.sleep(10)
                
            except Exception as e:
                print(f"Error during generation polling: {str(e)}")
                print(traceback.format_exc())
                time.sleep(10)
                continue
        
        # Download the video
        status_tracker.update_message("Downloading your masterpiece...", 0.95)
        try:
            response = requests.get(download_url, stream=True, timeout=30)
            response.raise_for_status()
            file_path = "output_video.mp4"
            with open(file_path, 'wb') as file:
                file.write(response.content)
            
            status_tracker.add_step("🎉 Video ready!", 1.0)
            return file_path, status_box
        except Exception as e:
            return None, f"📺 Video download failed: {str(e)}"
    
    except Exception as e:
        print(f"Error during generation: {str(e)}")
        print(traceback.format_exc())
        return None, f"🎪 The show must go on, but: {str(e)}"

# Create Gradio interface with a modern theme
with gr.Blocks(theme=gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="purple",
)) as app:
    gr.Markdown(
        """
        # 🎬 LumaAI Video Generator
        ### Transform your prompts into mesmerizing videos
        """
    )
    
    with gr.Row():
        # Create input and output columns
        prompt, camera_motion, api_key, image_input, generate_btn, status_display = create_input_column()
        video_output = create_output_column()
    
    generate_btn.click(
        fn=generate_video,
        inputs=[api_key, prompt, camera_motion, image_input],
        outputs=[video_output, status_display]
    )

if __name__ == "__main__":
    app.launch()