#!/usr/bin/env python3
"""Apply the site image treatment described in docs/image-brief.md.

Two jobs, both idempotent and re-runnable:

  cutout  remove the background from a product photograph      (needs rembg)
  treat   compose onto the dark canvas with the fade and glow  (needs Pillow)

Typical use:

  python3 tools/images/treat.py cutout raw/*.jpg      -> build/cut/*.png
  python3 tools/images/treat.py treat  build/cut/*.png -> content/assets/images/*.webp

If rembg is not available, cut the subject out by hand in GIMP, save as PNG with
transparency, and run only the `treat` step.
"""
import argparse, pathlib, sys

BG        = (0x0b, 0x0d, 0x10)   # --bg
GLOW      = (0x8e, 0xcb, 0xff)   # --link
SIZES     = {"hero": (1600, 1000), "detail": (1200, 1200)}
FADE_END  = 0.55                 # gradient reaches 0 opacity here
GLOW_A    = 0.32
GLOW_BLUR = 11
LIFT      = 1.06                 # shadow lift on the subject


def cutout(paths, outdir):
    try:
        from rembg import remove
    except ImportError:
        sys.exit("rembg not installed. See tools/images/README.md, or cut out by hand and skip this step.")
    from PIL import Image
    outdir.mkdir(parents=True, exist_ok=True)
    for p in paths:
        src = Image.open(p).convert("RGBA")
        dst = remove(src)
        target = outdir / (pathlib.Path(p).stem + ".png")
        dst.save(target)
        print(f"cut  {p} -> {target}")


def treat(paths, outdir, size="hero"):
    from PIL import Image, ImageFilter, ImageEnhance
    W, H = SIZES[size]
    outdir.mkdir(parents=True, exist_ok=True)

    for p in paths:
        subject = Image.open(p).convert("RGBA")

        # fit to 68% of canvas height, keep aspect
        scale = (H * 0.68) / subject.height
        subject = subject.resize((max(1, int(subject.width * scale)), int(H * 0.68)), Image.LANCZOS)

        # shadow lift
        rgb = ImageEnhance.Brightness(subject.convert("RGB")).enhance(LIFT)
        subject = Image.merge("RGBA", (*rgb.split(), subject.split()[3]))

        canvas = Image.new("RGB", (W, H), BG)

        # place right of centre
        x = int(W * 0.52)
        y = (H - subject.height) // 2
        canvas.paste(subject, (x, y), subject)

        # rim glow from the subject's own alpha
        alpha = subject.split()[3].filter(ImageFilter.GaussianBlur(GLOW_BLUR))
        glow = Image.new("RGB", subject.size, GLOW)
        canvas.paste(glow, (x, y), alpha.point(lambda a: int(a * GLOW_A / 255)))

        # left-to-right fade back to the page colour
        fade = Image.new("L", (W, 1))
        end = int(W * FADE_END)
        for px in range(W):
            fade.putpixel((px, 0), 255 if px < end else max(0, int(255 * (1 - (px - end) / max(1, W - end)))))
        fade = fade.resize((W, H))
        canvas = Image.composite(Image.new("RGB", (W, H), BG), canvas, fade)

        target = outdir / (pathlib.Path(p).stem + ".webp")
        canvas.save(target, "WEBP", quality=82, method=6)
        kb = target.stat().st_size / 1024
        print(f"treat {p} -> {target}  ({kb:.0f} KB)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("step", choices=["cutout", "treat"])
    ap.add_argument("paths", nargs="+", type=pathlib.Path)
    ap.add_argument("--out", type=pathlib.Path, default=None)
    ap.add_argument("--size", choices=list(SIZES), default="hero")
    a = ap.parse_args()
    out = a.out or pathlib.Path("build" if a.step == "cutout" else "content/assets/images")
    (cutout if a.step == "cutout" else treat)(a.paths, out) if a.step == "cutout" else treat(a.paths, out, a.size)


if __name__ == "__main__":
    main()