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

def remove_background(image_path, tolerance=30):
    try:
        img = Image.open(image_path).convert("RGBA")
        datas = img.getdata()
        
        # Sample the corner pixel to guess the "background" theme
        # But since it's a checkerboard, precise color matching fails.
        # We'll use a seed-fill (flood fill) approach from the corners.
        
        # Convert to numpy array for faster processing
        arr = np.array(img)
        
        # Create a mask for visited pixels (background)
        h, w = arr.shape[:2]
        mask = np.zeros((h, w), dtype=bool)
        
        # Stack for flood fill: (r, c)
        stack = [(0, 0), (0, w-1), (h-1, 0), (h-1, w-1)]
        
        # We need to be careful not to delete the object.
        # Assumption: The checkerboard is grey-scale or near grey-scale.
        # And the object (tree) is colorful (green/brown).
        
        # Helper to check if a pixel is "grey-ish" (background candidate)
        def is_background_candidate(pixel):
            r, g, b, a = pixel
            # Check for low saturation (grey/white/black)
            # max(r,g,b) - min(r,g,b) should be small for greys
            saturation = max(r, g, b) - min(r, g, b)
            return saturation < tolerance

        # New approach: smart flood fill with color tolerance is hard on noisy checkerboard.
        # Let's try a simpler heuristic first: 
        # Iterate all pixels. If a pixel is GREY-SCALE (within tolerance), make it transparent.
        # This risks deleting grey parts of the object (e.g. stones, bark).
        
        newData = []
        for item in datas:
            # item is (r, g, b, a)
            r, g, b, a = item
            
            # Check if it's a grey/white/black pixel
            # Checkerboard usually consists of light grey/white and dark grey blocks.
            # We enforce that R, G, and B are close to each other.
            if abs(r - g) < tolerance and abs(g - b) < tolerance and abs(r - b) < tolerance:
                # Also check brightness to avoid deleting dark black outlines if they are pure black
                # Let's say we only delete "light" greys/whites? 
                # The "dark" squares in the checkerboard from the log were around (77, 77, 77). 
                # This is quite dark.
                
                # If we just delete all greys, we might lose outlines (usually (0,0,0)).
                # So we check if it's NOT pure black.
                if r > 20 and g > 20 and b > 20: 
                     newData.append((0, 0, 0, 0)) # Transparent
                else:
                     newData.append(item)
            else:
                newData.append(item)

        img.putdata(newData)
        
        # Save back
        img.save(image_path, "PNG")
        print(f"Processed: {image_path}")
        return True
        
    except Exception as e:
        print(f"Error processing {image_path}: {e}")
        return False

# List of specific files to fix first (Trees)
targets = [
    "assets/references/trees/apple/apple_tree.png"
]

# Or scan directory
target_dir = "assets/references"

if __name__ == "__main__":
    # If arguments provided, use them
    if len(sys.argv) > 1:
        files = sys.argv[1:]
        for f in files:
             remove_background(f)
    else:
        # Recursive scan for testing
        for root, dirs, files in os.walk(target_dir):
            for file in files:
                if file.endswith("apple_tree.png"): # Safety: only target the known bad one first
                    full_path = os.path.join(root, file)
                    remove_background(full_path)