#!/usr/bin/env python3
"""Knock a solid background out of AI-generated sprites → transparent PNG.

Usage: tools/.venv/bin/python tools/dekey.py <in.jpeg> [more.jpeg ...]
Writes <name>.png next to each input. Key color is sampled from the corners
(works for Flow's magenta/white/gray isolation backgrounds), alpha is a soft
ramp on color distance, plus a despill pass so edges don't glow key-colored.
"""
import sys, pathlib
from PIL import Image

NEAR, FAR = 60, 130   # distance ramp: <NEAR fully transparent, >FAR fully opaque

def dekey(path):
    im = Image.open(path).convert('RGB')
    w, h = im.size
    px = im.load()
    # key = average of the four corners
    corners = [px[0, 0], px[w - 1, 0], px[0, h - 1], px[w - 1, h - 1]]
    key = tuple(sum(c[i] for c in corners) // 4 for i in range(3))
    out = Image.new('RGBA', (w, h))
    po = out.load()
    for y in range(h):
        for x in range(w):
            r, g, b = px[x, y]
            d = ((r - key[0]) ** 2 + (g - key[1]) ** 2 + (b - key[2]) ** 2) ** 0.5
            a = 0 if d < NEAR else 255 if d > FAR else int(255 * (d - NEAR) / (FAR - NEAR))
            if 0 < a < 255:   # despill: pull edge pixels away from the key color
                t = a / 255
                r = int(r * t + (r + g) / 2 * (1 - t))
                b = int(b * t + (b + g) / 2 * (1 - t))
            po[x, y] = (r, g, b, a)
    dst = pathlib.Path(path).with_suffix('.png')
    out.save(dst)
    print(f'{path} → {dst.name} (key rgb{key})')

for p in sys.argv[1:]:
    dekey(p)