#!/usr/bin/env python3
"""Paper-doll compositor for 90sDJsim (Lane C art).

Flux can't emit alpha, so garments are generated on flat white and keyed here.
One bare BASE body (704x1408); garments are keyed, scaled + placed onto the SAME
704x1408 canvas per web/world/doll_layout.json, so lane D just stacks them at (0,0)
in draw order: base -> bottom -> shoes -> top -> hat.

  key     IN.jpg OUT.png                 key one image (base body, staff face) to alpha
  install RAWDIR LAYOUT.json OUTDIR       key+scale+place every raw garment jpg -> 704x1408 png,
                                          then write the effective full layout back to LAYOUT.json
  fit     OVERLAYDIR OUT.png stem...      composite base.png + named overlays (draw order) -> preview
  sheet   OUT.png cols cellw f[:label]... grid contact sheet

stdlib + PIL/numpy only.  # ponytail: border-flood keyer over a naive white->alpha so
white garment details survive; placement lives in json so alignment is tunable without regen.
"""
import sys, json, os
import numpy as np
from PIL import Image, ImageDraw, ImageFilter, ImageFont

CANVAS = (704, 1408)  # contract §4 doll canvas

# per-slot default placement on the 704x1408 canvas: (center_x, top_y, target_width)
# measured off base.png (head 136-260 cx355, shoulders ~270, waist ~685, feet 1300-1345).
SLOT_DEF = {
    "hat":    (355,  90, 150),
    "top":    (352, 258, 330),
    "bottom": (352, 635, 265),
    "shoes":  (352, 1225, 225),
    "bag":    (313, 234, 419),   # cross-body record bag: strap tip ~ right shoulder, bag at left hip
}


def key_white(im, thresh=110, feather=1.1, extra_seeds=()):
    """Key the frame-connected white background to alpha via border flood-fill.
    Interior whites (a white logo, white shoe) survive — only bg reachable from the
    frame edge through near-white is removed. extra_seeds: (x,y) points to also flood
    from, for white regions enclosed by the art (e.g. inside a bag-strap triangle)."""
    im = im.convert("RGB")
    w, h = im.size
    work = im.copy()
    SENT = (255, 0, 254)
    seeds = [(0, 0), (w - 1, 0), (0, h - 1), (w - 1, h - 1),
             (w // 2, 0), (0, h // 2), (w - 1, h // 2)] + list(extra_seeds)
    for s in seeds:
        if min(work.getpixel(s)[:3]) >= 230:      # only flood from a background-white seed
            ImageDraw.floodfill(work, s, SENT, thresh=thresh)
    a = np.asarray(work)
    bg = np.all(a == np.array(SENT), axis=-1)
    alpha = np.where(bg, 0, 255).astype(np.uint8)
    am = Image.fromarray(alpha).filter(ImageFilter.MinFilter(3))   # erode 1px -> eat halo
    am = am.filter(ImageFilter.GaussianBlur(feather))              # soften edge
    out = im.convert("RGBA")
    out.putalpha(am)
    return out


def place(raw_rgba, slot, x=0, y=0, s=1.0):
    """Scale a keyed garment to its slot width and paste onto a fresh 704x1408 canvas.
    x,y,s from layout nudge the slot default."""
    bbox = raw_rgba.getbbox()
    g = raw_rgba.crop(bbox) if bbox else raw_rgba
    cx, top_y, tw = SLOT_DEF[slot]
    scale = (tw / g.width) * s
    g = g.resize((max(1, round(g.width * scale)), max(1, round(g.height * scale))), Image.LANCZOS)
    canvas = Image.new("RGBA", CANVAS, (0, 0, 0, 0))
    canvas.alpha_composite(g, (round(cx - g.width / 2 + x), round(top_y + y)))
    return canvas


def load_font(sz=15):
    try:
        return ImageFont.truetype("/System/Library/Fonts/Supplemental/Arial.ttf", sz)
    except Exception:
        return ImageFont.load_default()


def cmd_key(inp, out):
    key_white(Image.open(inp)).save(out)
    print("keyed", out)


def cmd_install(rawdir, layout_path, outdir):
    os.makedirs(outdir, exist_ok=True)
    layout = json.load(open(layout_path)) if os.path.exists(layout_path) else {}
    eff = {}
    for fn in sorted(os.listdir(rawdir)):
        if not fn.lower().endswith((".jpg", ".jpeg", ".png")):
            continue
        stem = os.path.splitext(fn)[0]
        slot = stem.split("_")[0]
        if slot not in SLOT_DEF:
            print("skip (no slot)", fn); continue
        n = {"x": 0, "y": 0, "s": 1.0, **layout.get(stem, {})}
        keyed = key_white(Image.open(os.path.join(rawdir, fn)))
        place(keyed, slot, n["x"], n["y"], n["s"]).save(os.path.join(outdir, stem + ".png"))
        eff[stem] = n
        print("placed", stem)
    json.dump(eff, open(layout_path, "w"), indent=2, sort_keys=True)
    print("wrote layout", layout_path, f"({len(eff)} garments)")


DRAW_ORDER = ["bottom", "shoes", "top", "hat", "bag"]  # bag is worn over everything


def cmd_tint(inp, out, hexcolor, sat_thresh=45):
    """Recolour the saturated body of a placed overlay to a target hex, keeping
    neutrals (white label panel, highlights, dark trim) and alpha untouched.
    One keeper geometry -> N colourways; V/S scale multiplicatively so shading survives."""
    im = Image.open(inp).convert("RGBA")
    r, g, b = (int(hexcolor.lstrip("#")[i:i + 2], 16) for i in (0, 2, 4))
    th, ts, tv = Image.new("RGB", (1, 1), (r, g, b)).convert("HSV").getpixel((0, 0))
    a = np.asarray(im)
    hsv = np.asarray(im.convert("RGB").convert("HSV")).astype(np.float32)
    H, S, V = hsv[..., 0], hsv[..., 1], hsv[..., 2]
    mask = (S > sat_thresh) & (a[..., 3] > 0)
    if not mask.any():
        raise SystemExit(f"tint: no saturated pixels above {sat_thresh} in {inp}")
    rs, rv = float(np.median(S[mask])), float(np.median(V[mask]))
    H[mask] = th
    S[mask] = np.clip(S[mask] * (max(ts, 1) / max(rs, 1)), 0, 255)
    V[mask] = np.clip(V[mask] * (max(tv, 1) / max(rv, 1)), 0, 255)
    rgb = Image.fromarray(np.stack([H, S, V], -1).astype(np.uint8), "HSV").convert("RGB")
    outim = Image.fromarray(np.dstack([np.asarray(rgb), a[..., 3]]), "RGBA")
    outim.save(out)
    print("tinted", out, hexcolor)


def cmd_fit(overlaydir, out, stems):
    base = Image.open(os.path.join(overlaydir, "base.png")).convert("RGBA")
    comp = Image.new("RGBA", base.size, (0, 0, 0, 0))
    comp.alpha_composite(base)
    for slot in DRAW_ORDER:
        for st in stems:
            if st.split("_")[0] == slot:
                p = os.path.join(overlaydir, st + ".png")
                if os.path.exists(p):
                    comp.alpha_composite(Image.open(p).convert("RGBA"))
    Image.alpha_composite(Image.new("RGBA", comp.size, (235, 235, 238, 255)), comp).convert("RGB").save(out)
    print("fit", out, stems)


def cmd_sheet(out, cols, cellw, items):
    pad, lab = 8, 22
    cells = []
    for a in items:
        path, label = a.split(":", 1) if ":" in a else (a, os.path.basename(a))
        im = Image.open(path).convert("RGBA")
        im = Image.alpha_composite(Image.new("RGBA", im.size, (235, 235, 238, 255)), im).convert("RGB")
        h = round(im.height * cellw / im.width)
        cells.append((im.resize((cellw, h)), label))
    cellh = max(im.height for im, _ in cells)
    rows = (len(cells) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * cellw + (cols + 1) * pad, rows * (cellh + lab) + (rows + 1) * pad), (40, 40, 46))
    d, f = ImageDraw.Draw(sheet), load_font()
    for i, (im, label) in enumerate(cells):
        r, c = divmod(i, cols)
        x, y = pad + c * (cellw + pad), pad + r * (cellh + lab + pad)
        sheet.paste(im, (x, y))
        d.text((x + 3, y + im.height + 3), label, fill=(230, 230, 235), font=f)
    sheet.save(out)
    print("sheet", out, sheet.size)


if __name__ == "__main__":
    c = sys.argv[1] if len(sys.argv) > 1 else ""
    if c == "key":  # key IN OUT [x,y ...extra flood seeds for enclosed white regions]
        seeds = [tuple(map(int, s.split(","))) for s in sys.argv[4:]]
        key_white(Image.open(sys.argv[2]), extra_seeds=seeds).save(sys.argv[3])
        print("keyed", sys.argv[3], f"(+{len(seeds)} seeds)" if seeds else "")
    elif c == "tint":  # tint IN.png OUT.png '#rrggbb' [sat_thresh]
        cmd_tint(sys.argv[2], sys.argv[3], sys.argv[4], int(sys.argv[5]) if len(sys.argv) > 5 else 45)
    elif c == "install":
        cmd_install(sys.argv[2], sys.argv[3], sys.argv[4])
    elif c == "fit":
        cmd_fit(sys.argv[2], sys.argv[3], sys.argv[4:])
    elif c == "sheet":
        cmd_sheet(sys.argv[2], int(sys.argv[3]), int(sys.argv[4]), sys.argv[5:])
    else:
        print(__doc__); sys.exit(1)