import os
import numpy as np
from PIL import Image

# MAPPING: New Uploads (Step 337) to Target Filenames
file_map = {
    "uploaded_image_0_1768981323417.png": "gumb_start.png",
    "uploaded_image_1_1768981323417.png": "merilec_zdravja.png",
    "uploaded_image_2_1768981323417.png": "okvir_zarjavel.png",
    "uploaded_image_3_1768981323417.jpg": "amnezija_maska.png"
}

# DIRECTORIES
upload_dir = "/Users/davidkotnik/.gemini/antigravity/brain/998d0b10-1733-4e5b-85ed-249b986ba9b3"
dest_dir = "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI"

def clean_and_restore(img_path, output_path):
    if not os.path.exists(img_path):
        print(f"Skipping missing: {img_path}")
        return

    try:
        img = Image.open(img_path).convert("RGBA")
        data = np.array(img)
        
        # Color Channels
        r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
        
        # --- CHECKERBOARD REMOVAL ALGORITHM ---
        
        # 1. Detect White/Light Squares (>230)
        mask_white = (r > 230) & (g > 230) & (b > 230)
        
        # 2. Detect Grey Squares (Neutral Grey, Range 160-225)
        is_neutral = (np.abs(r.astype(int) - g.astype(int)) < 15) & \
                     (np.abs(r.astype(int) - b.astype(int)) < 15)
        is_grey_range = (r > 160) & (r < 225)
        mask_grey = is_neutral & is_grey_range
        
        # 3. Detect Green Screen (Safety check)
        mask_green = (g > 200) & (r < 100) & (b < 100)
        
        # Combined Mask
        # Note: white_mask variable was accidentally scoped out or typoed in previous attempt? 
        # Actually it was defined above line 32. 
        # The error suggests execution order or scope issue. 
        # Let's redefine cleanly.
        
        remove_mask = (r > 230) & (g > 230) & (b > 230) | \
                      (is_neutral & is_grey_range) | \
                      (g > 200) & (r < 100) & (b < 100)
        
        # Apply Transparency
        data[remove_mask, 3] = 0
        
        # Save
        result = Image.fromarray(data)
        result.save(output_path)
        print(f"Restored & Cleaned: {output_path}")
        
    except Exception as e:
        print(f"Error: {e}")

# RUN
for u_name, t_name in file_map.items():
    src = os.path.join(upload_dir, u_name)
    dst = os.path.join(dest_dir, t_name)
    clean_and_restore(src, dst)