from PIL import Image, ImageColor
import os, zipfile, tempfile, shutil

def recolor_icons(zip_file, hex_color):
    target_rgb = ImageColor.getrgb(hex_color)  # Fixed: handles "red", "#ff0000", etc.
    
    # Rest of the function...

    # Set up temp directories
    temp_input_dir = tempfile.mkdtemp()
    temp_output_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
    temp_output_zip.close()

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

        # Process 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:  # If not transparent
                                pixels[x, y] = (*target_rgb, a)

                    image.save(filepath)

        # Create new 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 Recoloring Tool")
    gr.Markdown(
        "Upload a **ZIP** file of `.png` icons and choose a color to apply to all visible pixels.\n"
        "Transparency is preserved — only the visible parts are recolored."
    )

    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()