import collections
from PIL import Image

def remove_background(image_path, output_path):
    # Load image and convert to RGBA
    img = Image.open(image_path).convert("RGBA")
    width, height = img.size
    pixels = img.load()
    
    # Visited grid to prevent loops
    visited = [[False for _ in range(height)] for _ in range(width)]
    
    # Queue for BFS
    queue = collections.deque()
    
    # Helper to check if color matches checkerboard background
    def is_background_color(r, g, b):
        # White squares
        if r > 240 and g > 240 and b > 240:
            return True
        # Grey squares
        if 210 < r < 240 and 210 < g < 240 and 210 < b < 240:
            # Check if it's neutral grey
            if abs(r - g) < 5 and abs(g - b) < 5:
                return True
        return False

    # Initialize queue with all border pixels
    for x in range(width):
        # Top border
        r, g, b, a = pixels[x, 0]
        if is_background_color(r, g, b):
            queue.append((x, 0))
            visited[x][0] = True
            
        # Bottom border
        r, g, b, a = pixels[x, height - 1]
        if is_background_color(r, g, b):
            queue.append((x, height - 1))
            visited[x][height - 1] = True
            
    for y in range(height):
        # Left border
        r, g, b, a = pixels[0, y]
        if is_background_color(r, g, b):
            queue.append((0, y))
            visited[0][y] = True
            
        # Right border
        r, g, b, a = pixels[width - 1, y]
        if is_background_color(r, g, b):
            queue.append((width - 1, y))
            visited[width - 1][y] = True

    # Perform BFS
    directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]
    while queue:
        cx, cy = queue.popleft()
        
        # Set this background pixel to fully transparent
        pixels[cx, cy] = (0, 0, 0, 0)
        
        # Check neighbors
        for dx, dy in directions:
            nx, ny = cx + dx, cy + dy
            if 0 <= nx < width and 0 <= ny < height:
                if not visited[nx][ny]:
                    r, g, b, a = pixels[nx, ny]
                    if is_background_color(r, g, b):
                        visited[nx][ny] = True
                        queue.append((nx, ny))

    # Save the transparent image
    img.save(output_path, "PNG")
    print(f"Background successfully removed and saved to {output_path}")

if __name__ == "__main__":
    remove_background(
        "public/images/slider-courier-mask-purple.png",
        "public/images/slider-courier-mask-purple.png"
    )