import gradio as gr
import torch
import traceback
from trippy_video_generator import TrippyVideoGenerator
import time
import os

# Initialize the generator globally for better performance
try:
    generator = TrippyVideoGenerator()
    CUDA_AVAILABLE = torch.cuda.is_available()
except Exception as e:
    print(f"Error initializing generator: {str(e)}")
    CUDA_AVAILABLE = False
    generator = None

def check_system_status():
    """Check system requirements and return status message"""
    status_messages = []
    
    # Check CUDA
    if CUDA_AVAILABLE:
        gpu_name = torch.cuda.get_device_name(0)
        memory = torch.cuda.get_device_properties(0).total_memory / 1e9  # Convert to GB
        status_messages.append(f"✅ GPU detected: {gpu_name} ({memory:.1f} GB)")
    else:
        status_messages.append("❌ No CUDA-capable GPU detected. This app requires an NVIDIA GPU.")
    
    # Check model initialization
    if generator is not None:
        status_messages.append("✅ Model initialized successfully")
    else:
        status_messages.append("❌ Error initializing model")
    
    return "\n".join(status_messages)

def validate_prompt(prompt):
    """Validate the input prompt"""
    if not prompt or len(prompt.strip()) == 0:
        return False, "Prompt cannot be empty"
    if len(prompt) > 500:
        return False, "Prompt is too long (maximum 500 characters)"
    return True, "Valid prompt"

def generate_video(prompt, num_frames, fps, progress=gr.Progress()):
    """Generate video with progress updates and error handling"""
    try:
        # Validate inputs
        is_valid, message = validate_prompt(prompt)
        if not is_valid:
            return None, message

        if not CUDA_AVAILABLE:
            return None, "Error: No CUDA-capable GPU detected. This application requires an NVIDIA GPU."

        # Create output directory if it doesn't exist
        os.makedirs("outputs", exist_ok=True)
        output_path = f"outputs/video_{int(time.time())}.mp4"

        # Generate video with progress updates
        progress(0, desc="Initializing...")
        
        generator.generate_video(
            prompt=prompt,
            output_path=output_path,
            num_frames=num_frames,
            fps=fps
        )
        
        progress(1, desc="Video generation complete!")
        
        if os.path.exists(output_path):
            return output_path, "Video generated successfully!"
        else:
            return None, "Error: Video file was not created"

    except Exception as e:
        error_msg = f"Error during video generation: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
        return None, error_msg

def create_interface():
    """Create the Gradio interface"""
    # System status at startup
    system_status = check_system_status()
    
    with gr.Blocks(title="Trippy Video Generator") as interface:
        gr.Markdown("""
        # 🌈 Trippy Video Generator
        Generate mesmerizing, psychedelic videos from text descriptions using AI.
        """)
        
        # System status
        gr.Markdown(f"## System Status\n{system_status}")
        
        with gr.Row():
            with gr.Column():
                # Input components
                prompt_input = gr.Textbox(
                    label="Enter your prompt",
                    placeholder="Example: cosmic journey through fractal dimensions",
                    lines=3
                )
                
                with gr.Row():
                    num_frames = gr.Slider(
                        minimum=15,
                        maximum=60,
                        value=30,
                        step=1,
                        label="Number of Frames"
                    )
                    fps = gr.Slider(
                        minimum=3,
                        maximum=30,
                        value=15,
                        step=1,
                        label="Frames per Second"
                    )
                
                generate_btn = gr.Button("Generate Video 🎬", variant="primary")
                
            with gr.Column():
                # Output components
                video_output = gr.Video(label="Generated Video")
                status_output = gr.Textbox(
                    label="Status",
                    placeholder="Generation status will appear here...",
                    lines=3
                )
        
        # Example prompts
        gr.Examples(
            examples=[
                ["cosmic journey through fractal dimensions", 30, 15],
                ["psychedelic butterfly in neon forest", 30, 15],
                ["abstract dreams in liquid colors", 30, 15],
                ["interdimensional portal with swirling energies", 30, 15]
            ],
            inputs=[prompt_input, num_frames, fps],
            outputs=[video_output, status_output],
            fn=generate_video,
            cache_examples=True,
        )
        
        # Set up event handler
        generate_btn.click(
            fn=generate_video,
            inputs=[prompt_input, num_frames, fps],
            outputs=[video_output, status_output]
        )
        
        # Additional information
        gr.Markdown("""
        ## Tips for Best Results
        - Be descriptive in your prompts
        - Include words like "psychedelic", "surreal", "colorful" for more trippy effects
        - Experiment with different frame counts and FPS
        - Higher FPS = smoother video but longer generation time
        
        ## Known Limitations
        - Requires a CUDA-capable GPU
        - Generation time varies based on settings and hardware
        - May produce unexpected results with very abstract prompts
        
        ## Need Help?
        If you encounter any issues, check the status message for details.
        """)
    
    return interface

if _name_ == "_main_":
    interface = create_interface()
    interface.launch(
        share=True,
        enable_queue=True,
        max_threads=1
    )