#!/usr/bin/env python3
"""Lay rendered PNGs out in a labelled grid — the before/after tool.

`contact_sheet.py` renders many GLBs in one Blender scene; it cannot pair a render made *now*
against one committed in a previous round, which is exactly what a look round is judged on. This
composites finished images: no Blender, no 3D, just PIL.

    python3 pipeline/montage.py OUT.png --cols 2 --title "R41 -> R42" \\
        before.png:"R41 as shipped" after.png:"R42 dressed"

Each input is `PATH[:CAPTION]`. Tiles are scaled to a common width, captions are drawn under them,
and a title band is drawn on top when `--title` is given. `--pair-gap` widens the gutter between
column pairs so an A/B row reads as pairs rather than as a strip.
"""
import os, sys
from PIL import Image, ImageDraw, ImageFont

BG = (250, 249, 246)
FG = (26, 26, 30)
SUB = (90, 92, 100)


def font(sz, bold=False):
    for p in ("/System/Library/Fonts/Supplemental/Arial Bold.ttf" if bold else
              "/System/Library/Fonts/Supplemental/Arial.ttf",
              "/System/Library/Fonts/Helvetica.ttc",
              "/Library/Fonts/Arial.ttf"):
        try:
            return ImageFont.truetype(p, sz)
        except Exception:
            continue
    return ImageFont.load_default()


def main():
    av = sys.argv[1:]
    out = av[0]
    cols = int(av[av.index("--cols") + 1]) if "--cols" in av else 2
    title = av[av.index("--title") + 1] if "--title" in av else ""
    width = int(av[av.index("--width") + 1]) if "--width" in av else 620
    pair_gap = int(av[av.index("--pair-gap") + 1]) if "--pair-gap" in av else 0
    items = []
    skip = set()
    for i, a in enumerate(av):
        if a in ("--cols", "--title", "--width", "--pair-gap"):
            skip.add(i); skip.add(i + 1)
    for i, a in enumerate(av):
        if i == 0 or i in skip or a.startswith("--"):
            continue
        path, _, cap = a.partition(":")
        if not os.path.exists(path):
            print(f"MISSING {path}")
            continue
        items.append((path, cap))

    tiles = []
    for path, cap in items:
        im = Image.open(path).convert("RGB")
        h = int(im.height * width / im.width)
        tiles.append((im.resize((width, h), Image.LANCZOS), cap))
    if not tiles:
        print("nothing to montage")
        return 1

    # Captions WRAP to the tile width. The first take let a long caption run under its neighbour and
    # the two collided in the middle of the sheet, which is the one thing a proof shot may not do.
    f_cap = font(19)
    def wrap(text, px):
        words, lines, cur = text.split(), [], ""
        for w in words:
            t = (cur + " " + w).strip()
            if dr0.textlength(t, font=f_cap) <= px or not cur:
                cur = t
            else:
                lines.append(cur); cur = w
        if cur:
            lines.append(cur)
        return lines
    dr0 = ImageDraw.Draw(Image.new("RGB", (8, 8)))
    tiles = [(im, wrap(cap, width - 6) if cap else []) for im, cap in tiles]
    cap_lines = max((len(c) for _, c in tiles), default=0)
    cap_h = (8 + 23 * cap_lines) if cap_lines else 0
    rows = (len(tiles) + cols - 1) // cols
    row_h = [max(t.height for t, _ in tiles[r * cols:(r + 1) * cols]) + cap_h for r in range(rows)]
    pad, title_h = 16, (56 if title else 0)
    W = pad + cols * (width + pad) + (pair_gap * (cols // 2 - 1) if cols > 2 else 0)
    H = title_h + pad + sum(h + pad for h in row_h)
    sheet = Image.new("RGB", (W, H), BG)
    dr = ImageDraw.Draw(sheet)
    if title:
        dr.text((pad, 16), title, fill=FG, font=font(26, True))
    y = title_h + pad
    for r in range(rows):
        x = pad
        for c in range(cols):
            k = r * cols + c
            if k >= len(tiles):
                break
            im, cap = tiles[k]
            sheet.paste(im, (x, y))
            for li, line in enumerate(cap):
                dr.text((x + 2, y + im.height + 8 + li * 23), line, fill=SUB, font=f_cap)
            x += width + pad + (pair_gap if (pair_gap and c % 2 == 1) else 0)
        y += row_h[r] + pad
    os.makedirs(os.path.dirname(os.path.abspath(out)) or ".", exist_ok=True)
    sheet.save(out)
    print(f"MONTAGE {out}  {len(tiles)} tiles  {W}x{H}")
    return 0


if __name__ == "__main__":
    sys.exit(main())