import os
from PIL import Image, ImageDraw

def process_image(file_path, mode="flood_corners"):
    try:
        img = Image.open(file_path).convert("RGBA")
        width, height = img.size
        pixels = img.load()

        # Define checkerboard colors to target
        # Typically White and Gray. 
        # We'll treat anything very light/grey as background if it's connected to start point.
        # Gray is often (204, 204, 204) or (192, 192, 192)
        
        def is_checkerboard(p):
            r, g, b, a = p
            # Check for white-ish
            if r > 240 and g > 240 and b > 240: return True
            # Check for neutral grey-ish
            if abs(r - g) < 10 and abs(r - b) < 10 and 150 < r < 230: return True
            return False

        # Flood fill algorithm
        visited = set()
        queue = []

        if mode == "flood_corners":
            # Start from all 4 corners
            starts = [(0, 0), (width-1, 0), (0, height-1), (width-1, height-1)]
            for s in starts:
                if is_checkerboard(pixels[s]):
                    queue.append(s)
                    visited.add(s)

        elif mode == "flood_center":
            # For the mask, we want to clear the center
            center = (width // 2, height // 2)
            if is_checkerboard(pixels[center]):
                queue.append(center)
                visited.add(center)
                
        # Processing queue
        while queue:
            x, y = queue.pop(0)
            pixels[x, y] = (0, 0, 0, 0) # Make transparent

            # Check neighbors
            for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                nx, ny = x + dx, y + dy
                if 0 <= nx < width and 0 <= ny < height:
                    if (nx, ny) not in visited:
                        if is_checkerboard(pixels[nx, ny]):
                            visited.add((nx, ny))
                            queue.append((nx, ny))

        img.save(file_path)
        print(f"Processed: {file_path}")

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

# Paths
ui_dir = "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI"

# 1. Solid Objects (Frame, Meter, Button) - Flood from corners
solid_objects = [
    "okvir_zarjavel.png",
    "merilec_zdravja.png",
    "gumb_start.png"
]

for f in solid_objects:
    path = os.path.join(ui_dir, f)
    if os.path.exists(path):
        process_image(path, mode="flood_corners")

# 2. Mask (Vignette) - Flood from center (remove the middle checkerboard)
mask_path = os.path.join(ui_dir, "amnezija_maska.png")
if os.path.exists(mask_path):
    process_image(mask_path, mode="flood_center")