import gradio as gr
import config
from metrics import compute_mse
from mosaic_generator import create_mosaic
import time

def create_gradio_interface():
    """
    Create a Gradio interface for the mosaic generator.
    """

    def process_image(input_image, grid_size):
        """
        Process image and return mosaic.

        Parameters:
        -----------
        input_image : PIL.Image
            Input image from Gradio
        grid_size : int
            Grid size (will be grid_size × grid_size)

        Returns:
        --------
        tuple : (mosaic_image, mse_value, message)
        """
        try:
            # Update grid size
            config.grid_size = (grid_size, grid_size)

            # Create mosaic
            start = time.time()
            mosaic = create_mosaic(input_image, "tiles", config.grid_size)
            end = time.time()

            # Compute MSE
            mse = compute_mse(input_image, mosaic)

            message = f"Mosaic created successfully!\nGrid: {grid_size}×{grid_size}\nMSE: {mse:.2f}. \nThis image took {end-start} seconds to make!"
            return mosaic, message

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

    # Create Gradio interface
    with gr.Blocks(title="Image Mosaic Generator") as demo:
        gr.Markdown("# 🎨 Image Mosaic Generator")
        gr.Markdown("Upload an image and create a mosaic version using colored tiles!")

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(type="pil", label="Input Image")
                grid_size = gr.Slider(
                    minimum=8,
                    maximum=64,
                    value=32,
                    step=8,
                    label="Grid Size (cells per side)"
                )
                generate_btn = gr.Button("Generate Mosaic", variant="primary")

            with gr.Column():
                output_image = gr.Image(label="Mosaic Output")
                output_text = gr.Textbox(label="Status", lines=3)

        # Connect components
        generate_btn.click(
            fn=process_image,
            inputs=[input_image, grid_size],
            outputs=[output_image, output_text]
        )

        # Add examples if available
        gr.Markdown("### 💡 Tips")
        gr.Markdown("""
        - Start with smaller grid sizes (16×16) for faster processing
        - Larger grid sizes create more detailed mosaics but take longer
        - The default tiles are simple colored squares - you must add your own tiles to the 'tiles' folder!
        """)

    return demo


if __name__ == "__main__":
    # Create and launch the Gradio interface
    demo = create_gradio_interface()
    demo.launch()