import gradio as gr
from PIL import Image
import pillow_heif
import os
import shutil

# 1. Register HEIF opener with Pillow so it can handle .heic/.heif files
pillow_heif.register_heif_opener()

def format_bytes(size):
    # Helper to make file sizes readable
    power = 2**10
    n = 0
    power_labels = {0 : '', 1: 'KB', 2: 'MB', 3: 'GB', 4: 'TB'}
    while size > power:
        size /= power
        n += 1
    return f"{size:.2f} {power_labels[n]}"

def compress_image(image_path):
    if image_path is None:
        return None, "No image uploaded."

    try:
        # Get original file info
        original_size = os.path.getsize(image_path)
        filename = os.path.basename(image_path)
        name_without_ext = os.path.splitext(filename)[0]
        
        # Define output path
        output_filename = f"{name_without_ext}_compressed.png"
        output_path = os.path.join(os.path.dirname(image_path), output_filename)

        # Open the image
        # pillow-heif handles the HEIC/HEIF format automatically here
        with Image.open(image_path) as img:
            
            # Create a copy of the image to remove specific metadata that might 
            # conflict with saving (though we keep pixel data lossless)
            # We convert to RGBA if it supports transparency, or RGB otherwise
            if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
                img = img.convert('RGBA')
            else:
                img = img.convert('RGB')

            # SAVE AS PNG
            # compress_level=9 is maximum compression (slowest, smallest size)
            # optimize=True performs additional passes to reduce size
            img.save(output_path, "PNG", compress_level=9, optimize=True)

        # Get new file info
        new_size = os.path.getsize(output_path)
        
        # Calculate stats
        diff = new_size - original_size
        status = "Reduced" if diff < 0 else "Increased"
        percent = (abs(diff) / original_size) * 100
        
        info_text = (
            f"Original Size: {format_bytes(original_size)}\n"
            f"Compressed Size: {format_bytes(new_size)}\n"
            f"Difference: {status} by {format_bytes(abs(diff))} ({percent:.1f}%)\n\n"
            f"Note: Converting JPEGs/HEICs (lossy) to PNG (lossless) may increase file size "
            f"because PNG preserves the exact pixel data without compression artifacts."
        )

        return output_path, info_text

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

# 2. Create the Gradio Interface
with gr.Blocks(title="Lossless PNG Compressor") as app:
    gr.Markdown("## 🗜️ Lossless Image Compressor (Max PNG Level)")
    gr.Markdown(
        "Upload any image (JPG, PNG, HEIC, WEBP). It will be converted to **PNG** "
        "saved at **Compression Level 9** (Maximum). \n\n"
        "*Note: This preserves pixel quality perfectly. However, converting a lossy photo (JPG/HEIC) "
        "to lossless PNG often results in a larger file size.*"
    )

    with gr.Row():
        with gr.Column():
            # Input: Expects a file path
            image_input = gr.Image(type="filepath", label="Upload Image")
            submit_btn = gr.Button("Compress / Convert", variant="primary")
        
        with gr.Column():
            # Output: Returns a downloadable file and text stats
            file_output = gr.File(label="Download PNG")
            stats_output = gr.Textbox(label="Compression Stats")

    # Bind function
    submit_btn.click(
        fn=compress_image, 
        inputs=image_input, 
        outputs=[file_output, stats_output]
    )

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