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

def heic_to_png(heic_file, compression, lossless):
    # Read HEIC image
    heif_file = pillow_heif.read_heif(heic_file.name)
    image = Image.frombytes(
        heif_file.mode, heif_file.size, heif_file.data, "raw"
    )
    # Prepare output buffer
    png_bytes = io.BytesIO()
    if lossless:
        # Save with default PNG settings (lossless)
        image.save(png_bytes, format="PNG", optimize=True)
    else:
        # PNG compression: 0 (no compression, largest file) to 9 (max compression, smallest file)
        # Map slider 0-100 to PNG compress_level 0-9
        compress_level = int(9 - (compression / 100) * 9)
        image.save(
            png_bytes,
            format="PNG",
            compress_level=compress_level,
            optimize=True
        )
    png_bytes.seek(0)
    return png_bytes

with gr.Blocks() as demo:
    gr.Markdown("## HEIC to PNG Converter")
    gr.Markdown(
        "Upload a HEIC image, adjust compression, or enable 'Professional Lossless' for maximum PNG quality."
    )
    with gr.Row():
        heic_input = gr.File(label="Upload HEIC Image", file_types=[".heic"])
        png_output = gr.File(label="Converted PNG", file_types=[".png"])
    compression = gr.Slider(
        minimum=0,
        maximum=100,
        value=90,
        step=1,
        label="Compression Amount (0=Max compression, 100=Least compression)",
        info="Set to 100 for highest quality (least compression)."
    )
    lossless = gr.Checkbox(
        label="Professional Lossless (Ignore compression slider)",
        value=False
    )
    convert_btn = gr.Button("Convert")
    convert_btn.click(
        fn=heic_to_png,
        inputs=[heic_input, compression, lossless],
        outputs=png_output
    )

demo.launch()