import gradio as gr
from PIL import Image, ImageColor
import os
import zipfile
import tempfile
import shutil

def parse_color(color_str):
    """
    Parses a color string like "#ff0000", "red", or "rgba(255,0,0,1)" into an (R, G, B) tuple.
    """
    try:
        if color_str.startswith("rgba"):
            rgba = color_str.strip("rgba() ").split(",")
            r = int(float(rgba[0]))
            g = int(float(rgba[1]))
            b = int(float(rgba[2]))
            return (r, g, b)
        else:
            return ImageColor.getrgb(color_str)
    except Exception as e:
        raise ValueError(f"Invalid color input: {color_str} ({e})")

def recolor_icons(zip_file, selected_color):
    target_rgb = parse_color(selected_color)

    # Temp directories and output ZIP
    temp_input_dir = tempfile.mkdtemp()
    temp_output_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
    temp_output_zip.close()

    try:
        # Extract uploaded ZIP
        with zipfile.ZipFile(zip_file, 'r') as zip_ref:
            zip_ref.extractall(temp_input_dir)

        # Recolor PNGs
        for root, _, files in os.walk(temp_input_dir):
            for filename in files:
                if filename.lower().endswith('.png'):
                    filepath = os.path.join(root, filename)
                    image = Image.open(filepath).convert("RGBA")
                    pixels = image.load()

                    for y in range(image.height):
                        for x in range(image.width):
                            r, g, b, a = pixels[x, y]
                            if a != 0:  # not transparent
                                pixels[x, y] = (*target_rgb, a)

                    image.save(filepath)

        # Create output ZIP
        with zipfile.ZipFile(temp_output_zip.name, 'w') as zip_out:
            for root, _, files in os.walk(temp_input_dir):
                for file in files:
                    full_path = os.path.join(root, file)
                    arcname = os.path.relpath(full_path, temp_input_dir)
                    zip_out.write(full_path, arcname)

        return temp_output_zip.name

    finally:
        shutil.rmtree(temp_input_dir)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🎨 PNG Icon Recolor Tool")
    gr.Markdown(
        "Upload a **ZIP** file of `.png` images and pick any color.\n"
        "All non-transparent pixels will be recolored to your chosen color.\n"
        "Transparency is preserved. Get a new ZIP to download."
    )

    with gr.Row():
        zip_input = gr.File(label="📁 Upload ZIP of PNGs", file_types=[".zip"])
        color_picker = gr.ColorPicker(label="🎨 Choose Color", value="#ffffff")

    output_zip = gr.File(label="⬇️ Download Recolored ZIP")
    submit_btn = gr.Button("🔄 Recolor Icons")

    submit_btn.click(fn=recolor_icons, inputs=[zip_input, color_picker], outputs=output_zip)

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