import os
import numpy as np
from PIL import Image

# MAPPING: Uploaded File -> Target Name
file_map = {
    "uploaded_image_0_1768981117686.png": "gumb_start.png",
    "uploaded_image_1_1768981117686.png": "merilec_zdravja.png",
    "uploaded_image_2_1768981117686.png": "okvir_zarjavel.png",
    "uploaded_image_3_1768981117686.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"
if not os.path.exists(dest_dir):
    os.makedirs(dest_dir)

def clean_image(img_path, output_path, add_green_bg=False):
    try:
        if not os.path.exists(img_path):
            print(f"Skipping missing: {img_path}")
            return

        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 DETECTION LOGIC ---
        # 1. Pure White/Near White
        white_mask = (r > 230) & (g > 230) & (b > 230)
        
        # 2. Pure Grey/Near Grey
        # Standard checkerboard grey is often exactly equal R=G=B, around 204 or 192.
        # We look for low saturation (R~=G~=B) and specific brightness range.
        
        # Strict neutrality check (R, G, B within 10 of each other)
        is_neutral = (np.abs(r.astype(int) - g.astype(int)) < 15) & \
                     (np.abs(r.astype(int) - b.astype(int)) < 15)
                     
        # Brightness range for typical checkerboards (Light Grey)
        is_grey_range = (r > 160) & (r < 225)
        
        grey_mask = is_neutral & is_grey_range
        
        # 3. Green Screen (if any previous attempts left green)
        green_mask = (g > 200) & (r < 100) & (b < 100)
        
        # Combined Mask of "Pixels to Remove"
        remove_mask = white_mask | grey_mask | green_mask
        
        # --- APPLY CLEANING ---
        # Set Alpha to 0 where mask matches
        data[remove_mask, 3] = 0
        
        # --- OPTIONAL: GREEN BACKGROUND ---
        if add_green_bg:
            # Create a solid green image
            green_bg = np.zeros_like(data)
            green_bg[:,:] = [0, 255, 0, 255] # Solid Green
            
            # Where the original is NOT transparent, copy the original pixels
            # We use the INVERSE of the remove_mask (pixels we kept)
            # AND the original alpha channel (if original was significantly opaque)
            
            # Simplified: Composite original over green
            # But we just nuked the alpha of checkerboard, so we can standard composite
            
            cleaned_img = Image.fromarray(data)
            bg_img = Image.new("RGBA", cleaned_img.size, (0, 255, 0, 255))
            bg_img.alpha_composite(cleaned_img)
            
            # Save Green Version
            bg_img.save(output_path)
            print(f"Saved (GREEN BG): {output_path}")
            
        else:
            # Save Transparent Version
            processed_img = Image.fromarray(data)
            processed_img.save(output_path)
            print(f"Saved (TRANSPARENT): {output_path}")

    except Exception as e:
        print(f"Error processing {img_path}: {e}")

# EXECUTE
for u_file, target_name in file_map.items():
    src = os.path.join(upload_dir, u_file)
    
    # 1. Save Transparent Version (Standard game use)
    dst_trans = os.path.join(dest_dir, target_name)
    clean_image(src, dst_trans, add_green_bg=False)
    
    # 2. Save Green Version (As requested by user for manual keying?)
    # dst_green = os.path.join(dest_dir, "green_" + target_name)
    # clean_image(src, dst_green, add_green_bg=True)