import io
from PIL import Image
import gradio as gr

def get_file_size_label(size_bytes):
    if size_bytes < 1024 * 1024:
        return f"{size_bytes / 1024:.1f} KB"
    return f"{size_bytes / (1024 * 1024):.2f} MB"


def compress_to_target(img, fmt, target_bytes):
    buf = io.BytesIO()

    # ---- PNG (NO format change) ----
    if fmt == "PNG":
        for level in range(1, 10):
            buf = io.BytesIO()
            img.save(buf, format="PNG", optimize=True, compress_level=level)
            if buf.tell() <= target_bytes:
                buf.seek(0)
                return buf, "PNG"

        # If cannot reach target → return best effort PNG
        buf.seek(0)
        return buf, "PNG"

    # ---- JPEG ----
    elif fmt in ("JPEG", "JPG"):
        for q in range(95, 74, -2):
            buf = io.BytesIO()
            img.save(
                buf,
                format="JPEG",
                quality=q,
                optimize=True,
                subsampling=0 if q >= 85 else 2
            )
            if buf.tell() <= target_bytes:
                buf.seek(0)
                return buf, "JPEG"

        buf.seek(0)
        return buf, "JPEG"

    # ---- WEBP ----
    elif fmt == "WEBP":
        for q in range(95, 74, -2):
            buf = io.BytesIO()
            img.save(buf, format="WEBP", quality=q, method=6)
            if buf.tell() <= target_bytes:
                buf.seek(0)
                return buf, "WEBP"

        buf.seek(0)
        return buf, "WEBP"

    # ---- fallback ----
    else:
        img = img.convert("RGB")
        for q in range(95, 74, -2):
            buf = io.BytesIO()
            img.save(buf, format="JPEG", quality=q, optimize=True)
            if buf.tell() <= target_bytes:
                buf.seek(0)
                return buf, "JPEG"

        buf.seek(0)
        return buf, "JPEG"


def process_image(file, target_value, unit):
    if file is None:
        return None, "❌ Please upload an image"

    try:
        img = Image.open(file)
        fmt = (img.format or "JPEG").upper()

        # Normalize mode
        if img.mode in ("RGBA", "LA", "P"):
            if fmt == "JPEG":
                background = Image.new("RGB", img.size, (255, 255, 255))
                if img.mode == "P":
                    img = img.convert("RGBA")
                background.paste(
                    img,
                    mask=img.split()[-1] if img.mode in ("RGBA", "LA") else None
                )
                img = background
        elif img.mode not in ("RGB", "L"):
            img = img.convert("RGB")

        # Original size
        original_buf = io.BytesIO()
        if fmt == "PNG":
            img.save(original_buf, format="PNG")
        elif fmt == "WEBP":
            img.save(original_buf, format="WEBP")
        else:
            img.save(original_buf, format="JPEG", quality=95)

        original_size = original_buf.tell()

        target_bytes = (
            int(target_value * 1024)
            if unit == "KB"
            else int(target_value * 1024 * 1024)
        )

        # If already small
        if target_bytes >= original_size:
            return img, f"✅ Already optimized: {get_file_size_label(original_size)}"

        # Compress
        compressed_buf, out_fmt = compress_to_target(img, fmt, target_bytes)
        compressed_size = compressed_buf.getbuffer().nbytes

        saved = max(0, original_size - compressed_size)
        pct = (saved / original_size) * 100

        result_msg = (
            f"📦 Original: {get_file_size_label(original_size)}\n"
            f"🗜️ Compressed: {get_file_size_label(compressed_size)}\n"
            f"💾 Saved: {pct:.1f}%\n"
            f"📁 Format: {out_fmt}"
        )

        return Image.open(compressed_buf), result_msg

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


# ---- UI ----
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🗜️ ImagePress — Smart Compressor")
    gr.Markdown("✔ Preserves original format · ✔ No unwanted WEBP conversion")

    with gr.Row():
        input_img = gr.Image(type="filepath", label="Upload Image")
        output_img = gr.Image(label="Compressed Output")

    with gr.Row():
        target = gr.Number(value=200, label="Target Size")
        unit = gr.Radio(["KB", "MB"], value="KB", label="Unit")

    btn = gr.Button("⚡ Compress Image")

    result = gr.Textbox(label="Result")

    btn.click(
        fn=process_image,
        inputs=[input_img, target, unit],
        outputs=[output_img, result]
    )

demo.launch()