import os
import numpy as np
from PIL import Image
import shutil

# DIRECTORIES
upload_dir = "/Users/davidkotnik/.gemini/antigravity/brain/998d0b10-1733-4e5b-85ed-249b986ba9b3"
target_dir = "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI"
if not os.path.exists(target_dir):
    os.makedirs(target_dir)

# UPLOADED FILES
uploads = [
    "uploaded_image_0_1768959247545.png",
    "uploaded_image_1_1768959247545.png",
    "uploaded_image_2_1768959247545.png",
    "uploaded_image_3_1768959247545.jpg" # Note: .jpg might be the mask or frame?
]

def clean_checkerboard(img):
    # Convert manually to version with alpha
    img = img.convert("RGBA")
    data = np.array(img)
    
    r, g, b = data[:,:,0], data[:,:,1], data[:,:,2]
    
    # Define checkerboard colors (Broad range to catch compression artifacts)
    # White-ish
    mask_white = (r > 230) & (g > 230) & (b > 230)
    # Grey-ish
    mask_grey = (r > 190) & (r < 220) & \
                (g > 190) & (g < 220) & \
                (b > 190) & (b < 220) & \
                (np.abs(r.astype(int) - g.astype(int)) < 15)
                
    # Combine
    mask_to_remove = mask_white | mask_grey
    
    # Apply
    data[mask_to_remove, 3] = 0
    return Image.fromarray(data)

def identify_and_save(filename):
    src = os.path.join(upload_dir, filename)
    if not os.path.exists(src):
        print(f"Missing: {src}")
        return

    try:
        img = Image.open(src)
        w, h = img.size
        aspect = w / h
        
        name = "unknown.png"
        
        # LOGIC TO IDENTIFY IMAGES BY SIZE/ASPECT
        if w > 1500: # Full screen
            name = "amnezija_maska.png"
        elif aspect > 3.0: # Very wide -> Button or Frame
            if w < 600:
                name = "gumb_start.png" 
            else:
                name = "okvir_zarjavel.png"
        elif 0.8 < aspect < 1.2: # Square -> Gauge
            name = "merilec_zdravja.png"
        else:
            # Fallback for wide frame if logic fails
            if w > 600:
                 name = "okvir_zarjavel.png"
            else:
                 name = "gumb_start.png"
                 
        print(f"Identified {filename} ({w}x{h}) as {name}")
        
        # CLEAN
        cleaned_img = clean_checkerboard(img)
        
        # SAVE
        dst = os.path.join(target_dir, name)
        cleaned_img.save(dst)
        print(f"Saved to {dst}")
        
    except Exception as e:
        print(f"Error processing {filename}: {e}")

# EXECUTE
for u in uploads:
    identify_and_save(u)