import os
import numpy as np
from PIL import Image

# MAPPING
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"
}

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 process_image_green(u_name, t_name):
    src = os.path.join(upload_dir, u_name)
    dst = os.path.join(dest_dir, t_name)
    
    if not os.path.exists(src):
        print(f"MISSING SOURCE: {src}")
        return

    try:
        img = Image.open(src).convert("RGBA")
        data = np.array(img)
        
        r, g, b = data[:,:,0], data[:,:,1], data[:,:,2]
        
        # 1. CLEAN CHECKERBOARD
        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)
        
        # White & Grey & (Green safety)
        remove_mask = (r > 230) & (g > 230) & (b > 230) | \
                      (is_neutral & is_grey_range) | \
                      (g > 200) & (r < 100) & (b < 100)
                      
        data[remove_mask, 3] = 0
        
        cleaned = Image.fromarray(data)
        
        # 2. ADD GREEN BACKGROUND
        green_bg = Image.new("RGBA", cleaned.size, (0, 255, 0, 255))
        green_bg.alpha_composite(cleaned)
        
        # Save as RGB (No transparency, just Green Key)
        final = green_bg.convert("RGB")
        
        final.save(dst)
        print(f"CREATED: {dst}")
        
        # Verify
        if os.path.exists(dst):
             print(f" VERIFIED EXISTS: {os.path.getsize(dst)} bytes")
        else:
             print(f" ERROR: File not found after save!")

    except Exception as e:
        print(f"EXCEPTION: {e}")

print("STARTING RESTORE WITH GREEN BG...")
for u, t in file_map.items():
    process_image_green(u, t)
    
print("DIRECTORY LISTING:")
print(os.listdir(dest_dir))