#!/usr/bin/env python3
"""Composite transparent cut-outs (Kenerate remove_background results) onto brand colours, gradients or images, with a
soft contact shadow, consistent padding and any output size. Every cut-out x every background, plus a grid.

Examples
  python composite.py --cutouts https://.../a.png https://.../b.png --bg "#FF5A1F" --bg "gradient:#0F172A,#7C3AED" --size 1080x1080
  python composite.py --cutouts ./cutouts --bg "image:https://.../studio.jpg" --size 1080x1350 --align bottom --shadow soft
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import kmedia as K  # noqa: E402


def make_bg(spec: str, W: int, H: int, tmp: Path, st):
    s = spec.strip()
    if s.lower().startswith("image:"):
        return K.cover(K.open_image(K.fetch(s[6:], tmp, "bgimg", st)).convert("RGB"), W, H)
    bg = K.parse_bg(s, W, H)
    if bg is None:
        raise ValueError(f"Bad --bg '{spec}': use #RRGGBB, gradient:#a,#b or image:URL_OR_FILE")
    return bg


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--cutouts", nargs="+", required=True, help="transparent PNG URLs / files / a folder")
    ap.add_argument("--bg", action="append", default=[], help="#hex, gradient:#a,#b or image:URL (repeat)")
    ap.add_argument("--size", default="1080x1080")
    ap.add_argument("--padding", type=float, default=0.12, help="margin as a share of the frame (default 0.12)")
    ap.add_argument("--align", choices=["center", "bottom"], default="center")
    ap.add_argument("--shadow", choices=["soft", "hard", "none"], default="soft")
    ap.add_argument("--name", default="composite")
    ap.add_argument("--out-dir", default="kenerate-out/background-remove-batch")
    a = ap.parse_args()

    srcs = []
    for s in a.cutouts:
        p = Path(s)
        srcs += sorted(str(x) for x in p.iterdir() if x.suffix.lower() in (".png", ".webp")) if (not K.is_url(s) and p.is_dir()) else [s]
    bgs = a.bg or ["#FFFFFF"]
    st = K.Status(3, "background-remove-batch")
    K.require(ffmpeg=False, status=st)
    from PIL import Image, ImageFilter
    W, H = K.size_for(a.size, 1080)
    tmp = K.workdir()
    out = Path(a.out_dir)
    out.mkdir(parents=True, exist_ok=True)

    st.step(f"Fetching {len(srcs)} cut-out(s)")
    files = K.fetch_all(srcs, tmp, st, prefix="cut")
    backgrounds = [(b, make_bg(b, W, H, tmp, st)) for b in bgs]

    st.step(f"Compositing {len(files)} x {len(backgrounds)}")
    rows, made, k = [], [], 0
    for i, f in enumerate(files, 1):
        cut = K.open_image(f).convert("RGBA")
        alpha = cut.getchannel("A")
        if alpha.getextrema()[0] >= 250:
            st.warn(f"{f.name} has no transparency - run Kenerate remove_background on it first; used as-is")
        bb = alpha.point(lambda v: 255 if v > 10 else 0).getbbox() or (0, 0, cut.width, cut.height)
        cut = cut.crop(bb)
        bw, bh = W * (1 - 2 * a.padding), H * (1 - 2 * a.padding)
        s = min(bw / cut.width, bh / cut.height)
        cut = cut.resize((max(1, int(cut.width * s)), max(1, int(cut.height * s))), Image.LANCZOS)
        x = (W - cut.width) // 2
        y = (H - cut.height) // 2 if a.align == "center" else int(H * (1 - a.padding * 0.6)) - cut.height
        for j, (spec, bg) in enumerate(backgrounds, 1):
            canvas = bg.convert("RGBA").copy()
            if a.shadow != "none":
                sh = Image.new("RGBA", (W, H), (0, 0, 0, 0))
                ell = Image.new("L", (W, H), 0)
                from PIL import ImageDraw
                ew, eh = int(cut.width * 0.9), max(8, int(cut.height * 0.08))
                ImageDraw.Draw(ell).ellipse((x + (cut.width - ew) // 2, y + cut.height - eh // 2, x + (cut.width + ew) // 2,
                                             y + cut.height + eh // 2), fill=150 if a.shadow == "soft" else 210)
                ell = ell.filter(ImageFilter.GaussianBlur(max(4, eh // (2 if a.shadow == "soft" else 5))))
                sh.putalpha(ell)
                canvas.alpha_composite(sh)
            canvas.alpha_composite(cut, (x, y))
            dst = out / f"{K.safe_name(a.name)}_{i:02d}_bg{j}.png"
            canvas.convert("RGB").save(dst)
            made.append(dst)
            rows.append({"file": str(dst.resolve()), "cutout": f.name, "background": spec})
            k += 1
            st.progress(k * 100 / (len(files) * len(backgrounds)), dst.name)

    st.step("Grid")
    sheet = K.contact_sheet(made, out / f"{K.safe_name(a.name)}_grid.jpg", cols=len(backgrounds), cell=300,
                            labels=[r["background"][:24] for r in rows])
    (out / f"{K.safe_name(a.name)}_report.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    K.finish({"images": rows, "grid": str(sheet.resolve()), "size": f"{W}x{H}", "elapsed_s": st.elapsed()})


if __name__ == "__main__":
    K.main_guard(main)