import io
from pathlib import Path
from PIL import Image
import gradio as gr

# Optional support for extra formats requires plugins:
# - JPEG2000: built-in via Pillow
# - AVIF: install 'pillow-avif-plugin'
# - HEIF/HEIC: install 'pillow-heif'

# Supported formats and their default options
offered_formats = [
    "JPEG", "JPEG2000", "PNG", "BMP", "TIFF", "WEBP", "AVIF", "HEIF"
]

# Map formats to save kwargs defaults
format_options = {
    "JPEG": {"quality": 85, "optimize": True},
    "JPEG2000": {"quality_layers": [25, 50, 75], "codec": "jp2"},
    "PNG": {"compress_level": 6, "optimize": True},
    "BMP": {},
    "TIFF": {"compression": "tiff_deflate"},
    "WEBP": {"quality": 80, "lossless": False},
    "AVIF": {"quality": 50, "speed": 4},
    "HEIF": {"quality": 50, "subsampling": "4:2:0"}
}


def convert_image(
    img_file, output_format,
    quality, compress_level,
    jpeg2000_layers, avif_speed,
    subsampling, professional_lossless
):
    # img_file is a filesystem path (string)
    img_path = Path(img_file)
    img_bytes = img_path.read_bytes()

    # Load image and convert to RGB
    image = Image.open(io.BytesIO(img_bytes)).convert("RGB")

    # Prepare save kwargs
    fmt = output_format
    save_kwargs = {}

    if professional_lossless:
        # Force true lossless: use PNG
        fmt = "PNG"
        save_kwargs.update({"compress_level": 9, "optimize": True})
    else:
        # Base defaults for chosen format
        save_kwargs.update(format_options.get(fmt, {}))

        # Overrides per-format
        if fmt == "JPEG":
            save_kwargs["quality"] = quality
        elif fmt == "WEBP":
            save_kwargs["quality"] = quality
            save_kwargs["lossless"] = False
        elif fmt == "PNG":
            save_kwargs["compress_level"] = compress_level
            save_kwargs["optimize"] = True
        elif fmt == "JPEG2000":
            save_kwargs["quality_layers"] = [int(l) for l in jpeg2000_layers.split(",")]
        elif fmt == "AVIF":
            save_kwargs["quality"] = quality
            save_kwargs["speed"] = avif_speed
        elif fmt == "HEIF":
            save_kwargs["quality"] = quality
            save_kwargs["subsampling"] = subsampling

    # Save into a bytes buffer
    buf = io.BytesIO()
    image.save(buf, format=fmt, **save_kwargs)
    buf.seek(0)
    
    # Return bytes along with a filename
    filename = f"converted_image.{fmt.lower()}"
    return (buf.getvalue(), filename)


def build_app():
    with gr.Blocks() as demo:
        gr.Markdown("# Advanced Image Converter")

        with gr.Row():
            # Use 'filepath' so Gradio returns a local file path string
            inp = gr.File(label="Upload Image", file_types=["image"], type="filepath")
            out = gr.File(label="Download Converted")

        with gr.Row():
            fmt = gr.Dropdown(offered_formats, value="PNG", label="Output Format")
        with gr.Row():
            quality = gr.Slider(1, 100, value=85, step=1,
                                label="Quality (JPEG/AVIF/HEIF/WEBP)")
            compress = gr.Slider(0, 9, value=6, step=1,
                                 label="PNG Compress Level")
        with gr.Row():
            jpeg2000_layers = gr.Textbox(value="25,50,75",
                                         label="JPEG2000 Quality Layers (comma-separated)")
            avif_speed = gr.Slider(0, 8, value=4, step=1, label="AVIF Speed")
        subsample = gr.Dropdown(["4:4:4", "4:2:2", "4:2:0"], value="4:2:0",
                                label="HEIF Subsampling")
        professional = gr.Checkbox(label="Professional Lossless Mode",
                                   info="Force highest-quality lossless conversion as optimized PNG")
        convert_btn = gr.Button("Convert")

        convert_btn.click(
            fn=convert_image,
            inputs=[inp, fmt, quality, compress,
                    jpeg2000_layers, avif_speed,
                    subsample, professional],
            outputs=out
        )
    return demo


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