import gradio as gr
from PIL import Image
import os
from io import BytesIO

def get_gif_info(gif_file):
    """Get GIF dimensions and frame count."""
    try:
        with Image.open(gif_file) as img:
            width, height = img.size
            frames = 0
            try:
                while True:
                    img.seek(img.tell() + 1)
                    frames += 1
            except EOFError:
                img.seek(0)  # Reset to first frame
                return width, height, frames + 1, "Success"
    except Exception as e:
        return None, None, None, f"Error: {str(e)}"

def resize_gif(gif_file, width, height, lock_ratio):
    """Resize GIF while maintaining quality and aspect ratio."""
    try:
        # Open the GIF
        with Image.open(gif_file) as img:
            original_width, original_height = img.size
            frames = []
            
            # If lock_ratio is enabled, calculate height based on width
            if lock_ratio and width:
                height = int((width / original_width) * original_height)
            
            if not width or not height:
                return None, "Please provide valid width and/or height"
            
            # Process each frame
            try:
                while True:
                    # Convert frame to RGB if needed
                    frame = img.convert('RGBA')
                    # Resize frame using LANCZOS for better quality
                    resized_frame = frame.resize((int(width), int(height)), Image.LANCZOS)
                    frames.append(resized_frame)
                    img.seek(img.tell() + 1)
            except EOFError:
                pass

            # Create output BytesIO object
            output = BytesIO()
            
            # Save the resized GIF
            frames[0].save(
                output,
                format='GIF',
                save_all=True,
                append_images=frames[1:],
                loop=0,
                duration=img.info.get('duration', 100),
                optimize=True,
                quality=95
            )
            
            # Get new file size
            new_size = len(output.getvalue()) / 1024  # Size in KB
            
            # Reset BytesIO cursor and return
            output.seek(0)
            return output, f"Original Size: {original_width}x{original_height}\nNew Size: {width}x{height}\nFile Size: {new_size:.2f} KB"

    except Exception as e:
        return None, f"Error processing GIF: {str(e)}"

def update_height(width, original_width, original_height, lock_ratio):
    """Update height input when aspect ratio is locked."""
    if lock_ratio and width and original_width and original_height:
        try:
            return int((int(width) / original_width) * original_height)
        except:
            return None
    return None

def process_gif(gif_file, width, height, lock_ratio):
    """Main function to process the uploaded GIF."""
    if not gif_file:
        return None, "Please upload a GIF file", None
    
    # Get GIF info
    original_width, original_height, frame_count, error = get_gif_info(gif_file)
    
    if error:
        return None, error, None
    
    # Resize GIF
    output, message = resize_gif(gif_file, width, height, lock_ratio)
    
    return output, f"GIF Info:\nWidth: {original_width}px\nHeight: {original_height}px\nFrames: {frame_count}\n\n{message}", output

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# GIF Resizer")
    gr.Markdown("Upload a GIF, set the desired dimensions, and download the resized version.")
    msg = gr.Markdown()
    with gr.Row():
        with gr.Column():
            gif_input = gr.File(label="Upload GIF", file_types=[".gif"])
            lock_ratio = gr.Checkbox(label="Lock Aspect Ratio", value=True)
            
            with gr.Row():
                width_input = gr.Number(label="Width (px)", value=None)
                height_input = gr.Number(label="Height (px)", value=None)
            
            resize_button = gr.Button("Resize GIF")
        
        with gr.Column():
            output_info = gr.Textbox(label="GIF Information")
            output_gif = gr.Image(label="Resized GIF", type="filepath")
            download_button = gr.File(label="Download Resized GIF")
    
    # Event handlers
    gif_input.change(
        fn=get_gif_info,
        inputs=gif_input,
        outputs=[width_input, height_input, output_info, msg]
    )
    
    width_input.change(
        fn=update_height,
        inputs=[width_input, width_input, height_input, lock_ratio],
        outputs=height_input
    )
    
    resize_button.click(
        fn=process_gif,
        inputs=[gif_input, width_input, height_input, lock_ratio],
        outputs=[output_gif, output_info, download_button]
    )

demo.launch()