import sys
import os
import glob
from PIL import Image

def remove_background(img):
    img = img.convert("RGBA")
    datas = img.getdata()

    new_data = []
    for item in datas:
        r, g, b, a = item
        # 1. Near pure white
        if r > 240 and g > 240 and b > 240:
            new_data.append((255, 255, 255, 0))
        # 2. Near pure green
        elif g > 200 and g > r * 1.5 and g > b * 1.5:
            new_data.append((0, 255, 0, 0))
        # 3. Off-white/Light gray
        elif r > 220 and g > 220 and b > 220 and abs(r-g) < 15 and abs(g-b) < 15:
            new_data.append((r, g, b, 0))
        else:
            new_data.append(item)

    img.putdata(new_data)
    return img

def process_file(input_path, output_dir, tile_width=None, tile_height=None):
    try:
        img = Image.open(input_path)
        filename = os.path.basename(input_path).split('.')[0]
        
        if tile_width and tile_height:
            # Spritesheet mode
            w, h = img.size
            cols = w // tile_width
            rows = h // tile_height
            for r in range(rows):
                for c in range(cols):
                    left = c * tile_width
                    top = r * tile_height
                    right = left + tile_width
                    bottom = top + tile_height
                    
                    tile = img.crop((left, top, right, bottom))
                    tile = remove_background(tile)
                    
                    # Trim
                    bbox = tile.getbbox()
                    if bbox:
                        tile = tile.crop(bbox)
                        output_path = os.path.join(output_dir, f"{filename}_r{r}_c{c}.png")
                        tile.save(output_path, "PNG")
        else:
            # Single image mode
            img = remove_background(img)
            bbox = img.getbbox()
            if bbox:
                img = img.crop(bbox)
            output_path = os.path.join(output_dir, f"{filename}.png")
            img.save(output_path, "PNG")
            
        print(f"Processed: {input_path}")
    except Exception as e:
        print(f"Error processing {input_path}: {e}")

if __name__ == "__main__":
    # Usage: python process_asset.py <input_pattern> <output_dir> [tile_w tile_h]
    if len(sys.argv) < 3:
        print("Usage: python process_asset.py <input_pattern> <output_dir> [tile_w tile_h]")
    else:
        pattern = sys.argv[1]
        out_dir = sys.argv[2]
        os.makedirs(out_dir, exist_ok=True)
        
        tw = int(sys.argv[3]) if len(sys.argv) > 4 else None
        th = int(sys.argv[4]) if len(sys.argv) > 4 else None
        
        for f in glob.glob(pattern, recursive=True):
            if os.path.isfile(f):
                process_file(f, out_dir, tw, th)