#!/usr/bin/env python3
from __future__ import annotations
import argparse
from pathlib import Path
from PIL import Image, ImageFilter
import numpy as np

def fit_to_canvas(img: Image.Image, size=(1600, 900)) -> Image.Image:
    img = img.convert('RGBA')
    cw, ch = size
    ratio = min(cw / img.width, ch / img.height)
    nw, nh = max(1, int(round(img.width * ratio))), max(1, int(round(img.height * ratio)))
    res = img.resize((nw, nh), Image.LANCZOS)
    out = Image.new('RGBA', size, (0, 0, 0, 0))
    out.alpha_composite(res, ((cw - nw) // 2, (ch - nh) // 2))
    return out

def chroma_key_green(img: Image.Image) -> Image.Image:
    arr = np.asarray(img.convert('RGBA')).astype(np.float32)
    rgb = arr[..., :3]
    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]

    # Green dominance score: background is bright, saturated green; arrows/clothes are blue and should be preserved.
    max_rb = np.maximum(r, b)
    green_dom = g - max_rb
    green_ratio = g / (max_rb + 1.0)
    bright_green = (g > 95) & (green_dom > 28) & (green_ratio > 1.18)

    # Strong key where pure/near-pure green, softer on contact shadows.
    alpha = np.ones(g.shape, dtype=np.float32)
    strong = (g > 145) & (green_dom > 55) & (green_ratio > 1.35)
    soft = bright_green & ~strong
    alpha[strong] = 0.0
    # Soft region: preserve a little only if it is likely an actual shadow/edge, otherwise remove.
    shadow_like = soft & (g < 170) & (green_dom < 70)
    alpha[soft] = np.where(shadow_like[soft], 0.18, 0.0)

    # Preserve blue arrows and navy clothes aggressively.
    blue_keep = (b > 65) & (b > r * 1.08) & (b >= g * 0.75)
    dark_keep = (g < 95) | ((r > 80) & (g < 210))
    keep = blue_keep | dark_keep
    alpha[keep] = np.maximum(alpha[keep], 0.98)

    # Feather edge very slightly.
    a_img = Image.fromarray(np.clip(alpha * 255, 0, 255).astype(np.uint8), 'L')
    a_img = a_img.filter(ImageFilter.GaussianBlur(radius=0.45))
    alpha_blur = np.asarray(a_img).astype(np.float32) / 255.0
    alpha = np.maximum(alpha, alpha_blur * 0.72)
    alpha[strong] = 0.0

    # Despill: reduce green channel near semi-transparent edges and green-dominant pixels.
    edge = (alpha > 0.02) & (alpha < 0.98)
    greenish_edge = edge | ((alpha > 0.02) & bright_green)
    rgb[..., 1] = np.where(greenish_edge, np.minimum(g, np.maximum(r, b) + 18), g)

    out = np.dstack([rgb, alpha * 255.0]).clip(0, 255).astype(np.uint8)
    return Image.fromarray(out, 'RGBA')

def composite(img: Image.Image, color):
    bg = Image.new('RGBA', img.size, color + (255,))
    return Image.alpha_composite(bg, img)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--src', default='incoming/greenscreen-source')
    ap.add_argument('--dst', default='incoming/approved-figures')
    ap.add_argument('--qa', default='qa_screenshots')
    ap.add_argument('--only', nargs='*')
    args = ap.parse_args()
    src, dst, qa = Path(args.src), Path(args.dst), Path(args.qa)
    dst.mkdir(parents=True, exist_ok=True); qa.mkdir(parents=True, exist_ok=True)
    files = [p for p in sorted(src.glob('*')) if p.suffix.lower() in {'.png','.jpg','.jpeg','.webp'}]
    if args.only:
        want = set(args.only)
        files = [p for p in files if p.stem in want]
    for p in files:
        key = p.stem
        raw = Image.open(p).convert('RGBA')
        keyed = fit_to_canvas(chroma_key_green(raw), (1600, 900))
        out = dst / f'{key}.webp'
        keyed.save(out, 'WEBP', lossless=True, quality=100, method=6)
        composite(keyed, (0,0,0)).save(qa / f'{key}_qa_black.png')
        composite(keyed, (30,102,255)).save(qa / f'{key}_qa_blue.png')
        composite(keyed, (245,245,245)).save(qa / f'{key}_qa_light.png')
        print(f'KEYED {p.name} -> {out.name}')
if __name__ == '__main__':
    main()