import os
import gradio as gr
from PIL import Image
import torch
from diffusers import DiffusionPipeline
import tempfile

# Check for GPU availability
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

def initialize_model():
    """Initialize the Animator2D model."""
    try:
        # Initialize the pipeline
        pipeline = DiffusionPipeline.from_pretrained(
            "Lod34/Animator2D",
            trust_remote_code=True,
            device=DEVICE
        )
        return pipeline
    except Exception as e:
        raise Exception(f"Error initializing model: {str(e)}")

def generate_animation(
    description: str,
    action: str,
    direction: str,
    num_frames: int
):
    """Generate animation based on input parameters."""
    try:
        # Input validation
        if not all([description, action, direction]):
            raise ValueError("All text fields must be filled")
        
        # Initialize model
        pipeline = initialize_model()
        
        # Prepare prompt
        prompt = f"A sprite of {description} {action}, facing {direction}"
        
        # Generate animation
        output = pipeline(
            prompt=prompt,
            num_frames=num_frames,
            num_inference_steps=50
        )
        
        # Save animation as GIF
        temp_dir = tempfile.mkdtemp()
        output_path = os.path.join(temp_dir, "animation.gif")
        
        # Convert output frames to GIF
        frames = [Image.fromarray(frame) for frame in output.frames]
        frames[0].save(
            output_path,
            save_all=True,
            append_images=frames[1:],
            duration=100,
            loop=0
        )
        
        return output_path
        
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")

def create_interface():
    """Create and launch the Gradio interface."""
    
    with gr.Blocks(title="Animator2D Sprite Generator") as interface:
        gr.Markdown("# Animator2D Sprite Generator")
        gr.Markdown("Generate animated sprites using AI")
        
        with gr.Row():
            with gr.Column():
                # Input components
                description = gr.Textbox(
                    label="Sprite Description",
                    placeholder="E.g., a cute pixel art cat"
                )
                action = gr.Textbox(
                    label="Sprite Action",
                    placeholder="E.g., walking, jumping"
                )
                direction = gr.Dropdown(
                    label="Direction",
                    choices=["North", "South", "East", "West"],
                    value="South"
                )
                num_frames = gr.Slider(
                    label="Number of Frames",
                    minimum=2,
                    maximum=24,
                    value=8,
                    step=1
                )
                generate_btn = gr.Button("Generate Animation")
            
            with gr.Column():
                # Output components
                output_image = gr.Image(label="Generated Animation", type="filepath")
        
        # Connect components
        generate_btn.click(
            fn=generate_animation,
            inputs=[description, action, direction, num_frames],
            outputs=output_image
        )
        
    return interface

# Launch the application
if __name__ == "__main__":
    interface = create_interface()
    interface.launch(share=True)