import os
from PIL import Image

def remove_background(image_path, color_to_remove=None, tolerance=30, mode="top_left"):
    try:
        img = Image.open(image_path).convert("RGBA")
        datas = img.getdata()
        
        # Determine background color to remove
        bg_color = color_to_remove
        if bg_color is None:
            if mode == "top_left":
                bg_color = datas[0][:3] # Pick top-left pixel
            else:
                bg_color = (0,0,0) # Default placeholder to avoid NoneType error in zip
                
        newData = []
        for item in datas:
            # Check difference
            pixel = item[:3]
            
            # Simple Euclidean distance approximation or exact match
            # Let's do a strict match for green if specified, or threshold for checkerboard
            
            if mode == "green":
                # Check for Green (0, 255, 0) logic
                # Allow some tolerance for compression artifacts
                if item[1] > 200 and item[0] < 100 and item[2] < 100:
                    newData.append((255, 255, 255, 0))
                else:
                    newData.append(item)
            
            elif mode == "checkerboard":
                # Heuristic: if pixel is grey/white characteristic of checkerboard
                # Checkerboard usually alternates white (255) and grey (204 or similar)
                r, g, b = pixel
                if (r > 200 and g > 200 and b > 200) or (150 < r < 210 and 150 < g < 210 and 150 < b < 210):
                     # Likely checkerboard, but dragging risk of deleting white parts of image. 
                     # Using top-left reference is safer if image has borders.
                     diff = sum([abs(c1 - c2) for c1, c2 in zip(pixel, bg_color)])
                     if diff < tolerance:
                         newData.append((255, 255, 255, 0))
                     else:
                         newData.append(item)
                else:
                    newData.append(item)
                    
            else:
                newData.append(item)

        img.putdata(newData)
        img.save(image_path, "PNG")
        print(f"cleaned: {image_path}")
        
    except Exception as e:
        print(f"Error cleaning {image_path}: {e}")

# Paths
ui_path = "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI"
char_path = "/Users/davidkotnik/repos/novafarma/assets/slike/characters"
ghost_path = "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/Characters/starsa/Ghost"

# 1. Clean UI (Checkerboard)
# Note: The checkerboard seems to be distinct. Let's try to target the grey/white pattern. 
# Actually, since these were generated and saved with checkerboard 'baked in', it's tricky.
# I will aggressively target the specific checkerboard colors found in common generators.

# Or simpler: The user wants me to fix the scene. Cleaning assets is the permanent fix.

ui_files = [
    "okvir_zarjavel.png",
    "merilec_zdravja.png",
    "amnezija_maska.png"
]

for f in ui_files:
    full_path = os.path.join(ui_path, f)
    if os.path.exists(full_path):
        remove_background(full_path, mode="checkerboard", tolerance=50)

# 2. Clean Kai (Green Screen)
kai_path = os.path.join(char_path, "liki_kai_ref_kai.png")
if os.path.exists(kai_path):
    remove_background(kai_path, mode="green")

# 3. Clean Ghosts (Cyan/Green Screen?)
# The user mentioned green key for everyone, so I'll apply green filter.
ghosts = [
    "ghost_otac_cyan.png",
    "MOJE_SLIKE_KONCNA_ostalo_parents_transparent_ghosts_clean.png"
]
for g in ghosts:
    full_path = os.path.join(ghost_path, g)
    if os.path.exists(full_path):
        remove_background(full_path, mode="green")