import gradio as gr
from PIL import Image
import tempfile
import os
from image_processor import process_image, save_image_with_format
import io

def apply_standard_settings(setting):
    """Returns the parameters for the selected standard setting."""
    settings_dict = {
        "S light": (True, True, "240x240", 48, "whitesmoke"),
        "M light": (True, True, "480x480", 96, "whitesmoke"),
        "L light": (True, True, "960x960", 128, "whitesmoke"),
        "S dark": (True, True, "240x240", 48, "#2A373D"),
        "M dark": (True, True, "480x480", 96, "#2A373D"),
        "L dark": (True, True, "960x960", 128, "#2A373D"),
    }
    return settings_dict.get(setting, (None, None, None, None, None))

def settings_description(crop, remove_bg, resize, padding, background, quality):
    """Generate an HTML text description of the current settings."""
    description = f"""
    <ul style="font-size:small;">
        <li>Crop: {crop}</li>
        <li>Remove Background: {remove_bg}</li>
        <li>Resize: {resize if resize else 'No resize'}</li>
        <li>Padding: {padding}</li>
        <li>Background: {background}</li>
        <li>Quality: {quality}%</li>
    </ul>
    """
    return description


def get_file_info(file):
    """Get file info from the uploaded image file."""
    if file is None:
        return "No file uploaded."
    
    try:
        # For type="image", file is a dict with keys 'path' and 'name'
        if isinstance(file, dict):
            original_name = os.path.splitext(os.path.basename(file["name"]))[0]
            file_path = file["path"]
        else:
            original_name = os.path.splitext(os.path.basename(file.name))[0]
            file_path = file.name
        
        # Open the image file
        image = Image.open(file_path)
        
        # Get image dimensions
        width, height = image.size

        # Calculate file size
        size_bytes = os.path.getsize(file_path)
        size_kb = size_bytes / 1024

        return f"""<div style="text-align: left; padding: 5px;">
            <p>Original Name: {original_name}</p>
            <p>Size: {width}x{height}px ({size_kb:.1f} KB)</p>
        </div>"""
    except Exception as e:
        return f"Error getting file info: {str(e)}"


def gradio_interface(file, standard_settings, crop=False, remove_bg=False, 
                   resize=None, padding=0, background="white", 
                   quality=90):
    """Main interface function for the Gradio app."""
    quality = int(quality)

    if file is None:
        standard_image_path = './data/examples/supermario.png'
        file = standard_image_path
        original_name = os.path.splitext(os.path.basename(standard_image_path))[0]
    else:
        # Handle both dictionary (type="image") and file object cases
        if isinstance(file, dict):
            original_name = os.path.splitext(os.path.basename(file["name"]))[0]
            file_path = file["path"]
        else:
            original_name = os.path.splitext(os.path.basename(file.name))[0]
            file_path = file.name

    # Open the image file
    image = Image.open(file_path if isinstance(file, dict) else file)

    if standard_settings and standard_settings != "None":
        crop, remove_bg, resize, padding, background = apply_standard_settings(standard_settings)

    resize_dimensions = None
    if resize:
        try:
            width, height = map(int, resize.split('x'))
            resize_dimensions = (width, height)
        except ValueError:
            return None, None, None, None, "Invalid format for resize dimensions. Please use 'WxH'."

    # Process the image
    processed_image = process_image(
        image, 
        crop=crop, 
        remove_bg=remove_bg, 
        resize=resize_dimensions, 
        padding=padding, 
        background=background
    )

    # For preview, ensure we're working with RGB mode
    preview_image = processed_image
    if processed_image.mode in ('RGBA', 'LA'):
        preview_image = Image.new('RGB', processed_image.size, 'white')
        preview_image.paste(processed_image, mask=processed_image.split()[-1])

    # Get temp directory for saving files
    temp_dir = tempfile.gettempdir()

    # Save in multiple formats
    downloads = {}
    # Save WebP
    downloads['webp'] = save_image_with_format(
        processed_image,
        os.path.join(temp_dir, 'temp'),
        format='webp',
        quality=quality,
        custom_filename=original_name
    )

    # Save PNG
    downloads['png'] = save_image_with_format(
        processed_image,
        os.path.join(temp_dir, 'temp'),
        format='png',
        custom_filename=original_name
    )

    # Save JPG
    downloads['jpg'] = save_image_with_format(
        processed_image,
        os.path.join(temp_dir, 'temp'),
        format='jpg',
        quality=quality,
        custom_filename=original_name
    )

    # Generate settings description
    applied_settings = settings_description(
        crop, remove_bg, resize, padding, background, quality
    )

    # Return the preview image and downloads
    return (preview_image,  # Return the actual image object for preview
            downloads['webp'], 
            downloads['png'], 
            downloads['jpg'], 
            applied_settings)


# Define example images first (update to remove output_format parameter)
example_images = [
    [os.path.join("data", "examples", "supermario.png"),
     "S light", True, True, "480x420", 10, "whitesmoke", 90],
    [os.path.join("data", "examples",
                  "depositphotos_520707962-stock-photo-fujifilm-s10-body-black-fujifilm.jpg"),
     "None", True, True, "480x320", 48, "blue", 90],
    [os.path.join("data", "examples", "batman_b_c_320x280_bg.png"),
     "None", True, True, "360x360", 48, "yellow", 90],
]


def preview_and_info(f):
    """Handle file preview and info, including None case"""
    if f is None:
        return None, "No file uploaded."
    try:
        return Image.open(f.name), get_file_info(f)
    except Exception as e:
        return None, f"Error loading image: {str(e)}"
    

with gr.Blocks(title="IMAGER ___ v0.2 Image Processing Tool") as demo:
    gr.Markdown("# 🖼️ IMAGER v0.2 - Image Processing Tool 🛠️")
    gr.Markdown("📤 Upload an image and select one of the processing options below to get started! 🚀")
    
    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            # Image Upload Section
            with gr.Group():
                input_image = gr.Image(type="pil", label="Image Preview", scale=2)
                input_file = gr.File(label="Upload Image", scale=1)
                input_info = gr.HTML(label="Image Information")
            
            # Basic Settings Section
            with gr.Group():
                gr.Markdown("### Settings")
                settings = gr.Radio(
                    choices=["None", "S light", "M light", "L light", "S dark", "M dark", "L dark"],
                    label="Preset Settings",
                    value="None",
                    container=True
                )
                with gr.Row():
                    crop = gr.Checkbox(label="Crop on Object", container=True)
                    remove_bg = gr.Checkbox(label="Remove Background", container=True)
                
                # Size and Background Settings
                with gr.Row():
                    resize = gr.Textbox(
                        label="Resize (WxH)", 
                        placeholder="Example: 100x100",
                        container=True
                    )
                    padding = gr.Slider(
                        minimum=0, 
                        maximum=200, 
                        label="Padding",
                        container=True
                    )
                background = gr.Textbox(
                    label="Background", 
                    placeholder="Color name or hex code",
                    container=True
                )
                quality = gr.Slider(
                    minimum=1,
                    maximum=95,
                    value=85,
                    step=1,
                    label="Image Quality %",
                    container=True
                )

        # Output Section
        with gr.Column(scale=1):
            with gr.Group():
                gr.Markdown("### Preview and Downloads")
                output_image = gr.Image(type="pil", label="Processed Image Preview")
                settings_info = gr.HTML(label="Applied Settings")
                
                with gr.Row():
                    webp_download = gr.File(label="WebP")
                    png_download = gr.File(label="PNG")
                    jpg_download = gr.File(label="JPG")

    # Process Button
    with gr.Row():
        process_btn = gr.Button("Process Image", size="lg", variant="primary")

    # Examples Section
    with gr.Accordion("Examples", open=False):
        gr.Examples(
            examples=example_images,
            inputs=[
                input_file, settings, crop, remove_bg, resize, 
                padding, background, quality
            ]
        )

    # Wire up the events
    input_file.change(
        fn=preview_and_info,
        inputs=[input_file],
        outputs=[input_image, input_info]
    )
    
    process_btn.click(
        fn=gradio_interface,
        inputs=[
            input_file, settings, crop, remove_bg, resize, 
            padding, background, quality
        ],
        outputs=[
            output_image, webp_download, png_download, jpg_download, settings_info
        ]
    )

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