from rembg import remove, new_session
from PIL import Image
import os
import glob

# Paths to clean (Recursive search for crops)
SEARCH_PATTERNS = [
    "assets/PHASE_PACKS/*/crops/**/*.png",
    "assets/crops/**/*.png",
    "assets/references/trees/dead/dead_tree.png" # Also clean the tree while we're at it
]

# Initialize session for speed
session = new_session()

def clean_image(path):
    try:
        print(f"🌿 Processing: {path}")
        img = Image.open(path)
        
        # Use simple removal first (rabbit method)
        # If image is already transparent, rembg might do weird things, but usually safe.
        output = remove(img, session=session, alpha_matting=True)
        
        # Save inplace
        output.save(path)
        print(f"✨ Cleaned: {os.path.basename(path)}")
        
    except Exception as e:
        print(f"❌ Error cleaning {path}: {e}")

if __name__ == "__main__":
    print("🐇 The Rabbit Method: Cleaning all plants...")
    files = []
    for pattern in SEARCH_PATTERNS:
        found = glob.glob(pattern, recursive=True)
        files.extend(found)
        
    print(f"Found {len(files)} plant assets.")
    
    for f in files:
        # Skip if it's not a file
        if not os.path.isfile(f): continue
        if "_sheet" in f: continue # Skip spritesheets if any
        
        clean_image(f)
        
    print("✅ All plants cleaned!")