import gradio as gr
from PIL import Image
import io
import tempfile

def compress_image(input_image, output_format, quality=85):
    img_byte_arr = io.BytesIO()
    if output_format.upper() in ['JPEG', 'JPG']:
        input_image.save(img_byte_arr, format='JPEG', quality=quality)
    elif output_format.upper() == 'PNG':
        input_image.save(img_byte_arr, format='PNG', compress_level=int((100-quality)/10))
    elif output_format.upper() == 'WEBP':
        input_image.save(img_byte_arr, format='WEBP', quality=quality)
    
    img_byte_arr = img_byte_arr.getvalue()
    with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{output_format.lower()}') as temp_file:
        temp_file.write(img_byte_arr)
        temp_file_path = temp_file.name
    
    return temp_file_path

def image_compression_interface(uploaded_image, output_format, quality):
    try:
        image = Image.open(uploaded_image)
        if image.mode in ("RGBA", "P") and output_format.upper() in ['JPEG', 'JPG', 'WEBP']:
            image = image.convert("RGB")
        
        compressed_image_path = compress_image(image, output_format, quality)
        return compressed_image_path
    
    except Exception as e:
        return str(e)

def get_current_size(compressed_image_path):
    try:
        with open(compressed_image_path, 'rb') as f:
            size = len(f.read())
        return f"Current Size: {size} bytes"
    except FileNotFoundError:
        return "Image not found"

def update_ui(uploaded_image, output_format, quality):
    try:
        image = Image.open(uploaded_image)
        if image.mode in ("RGBA", "P") and output_format.upper() in ['JPEG', 'JPG', 'WEBP']:
            image = image.convert("RGB")
        
        compressed_image_path = compress_image(image, output_format, quality)
        current_size = get_current_size(compressed_image_path)
        return [compressed_image_path, current_size]
    
    except Exception as e:
        return [str(e), None]

def create_gradio_interface():
    image_input = gr.File(label="Upload Image")
    output_format_input = gr.Dropdown(choices=["JPEG", "JPG", "PNG", "WEBP"], label="Select Output Format")
    quality_input = gr.Slider(minimum=1, maximum=100, step=1, value=85, label="Select Compression Quality")
    size_display = gr.Textbox(label="Current Size")

    def on_quality_change(quality):
        try:
            image = Image.open(uploaded_image.value)
            if image.mode in ("RGBA", "P") and output_format_input.value.upper() in ['JPEG', 'JPG', 'WEBP']:
                image = image.convert("RGB")
            
            compressed_image_path = compress_image(image, output_format_input.value, quality)
            current_size = get_current_size(compressed_image_path)
            return current_size
        except Exception as e:
            return str(e)

    quality_input.change(on_quality_change, inputs=[quality_input], outputs=size_display)

    gr.Interface(
        fn=update_ui,
        inputs=[image_input, output_format_input, quality_input],
        outputs=[gr.File(label="Download Compressed Image"), size_display],
    ).launch()

if __name__ == "__main__":
    create_gradio_interface()