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

def compress_image(input_image, output_format, quality=85):
    """
    Compress an image and return the compressed image as a file.
    Parameters:
    - input_image (PIL.Image.Image): The input image.
    - output_format (str): The format of the output image ('JPEG', 'JPG', 'PNG', or 'WEBP').
    - quality (int, optional): The quality of the output image (1-100). Default is 85.
    Returns:
    - str: The path to the temporary file containing the compressed image.
    """
    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)
    
    # Save the compressed image to a temporary file
    img_byte_arr = img_byte_arr.getvalue()  # Convert to bytes
    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  # Get the temporary file path

    return temp_file_path

def image_compression_interface(uploaded_image, output_format, quality):
    try:
        # Open the image using PIL
        image = Image.open(uploaded_image)

        # Convert image to RGB if it has an alpha channel and the output format is JPEG, JPG, or WEBP
        if image.mode in ("RGBA", "P") and output_format.upper() in ['JPEG', 'JPG', 'WEBP']:
            image = image.convert("RGB")
        
        # Compress the image and save to a temporary file
        compressed_image_path = compress_image(image, output_format, quality)

        # Return the path to the compressed image file
        return compressed_image_path
    
    except Exception as e:
        return str(e)

# Create the Gradio interface
def create_gradio_interface():
    # Input components for Gradio (Updated for Gradio v3+)
    image_input = gr.File(label="Upload Image")  # Changed to gr.File to handle raw file upload
    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")

    # Output components for Gradio (Updated for Gradio v3+)
    image_output = gr.File(label="Download Compressed Image")  # Use gr.File for download

    # Set up the Gradio interface
    gr.Interface(
        fn=image_compression_interface,
        inputs=[image_input, output_format_input, quality_input],
        outputs=[image_output],
    ).launch()

if __name__ == "__main__":
    create_gradio_interface()