#!/usr/bin/env python3
"""Element-head teacher data, step 1: contact sheets showing each page image every way the extension can render it.

For each sampled image (logo, icon, formula, photo, backdrop...) of a crawled light page that has a Haiku dark design,
one row: the element in its light page (red outline), then the candidates on the dark surface designed for that spot:
  A keep    - unchanged
  B invert  - exactly Chrome's `invert(1) hue-rotate(180deg)` (dark <-> light, hues kept)
  C adapt   - the runtime's adaptPixels: neutral dark ink -> light, white box removed, brand colours kept
  D remove  - large CSS background images only: the designed dark surface shows through
Writes OUT/sheets/<id>.png and OUT/meta.<worker>.jsonl (one line per sheet: rows with page id, image index, offers).
"""
import argparse, base64, glob, gzip, hashlib, io, json, math, os, random, sys
from multiprocessing import Pool
from pathlib import Path

import numpy as np
from PIL import Image, ImageDraw, ImageFont

HUE180 = np.array([[-0.574, 1.43, 0.144], [0.426, 0.43, 0.144], [0.426, 1.43, -0.856]])
SMALL, LARGE = "small", "large"
GEOM = {SMALL: dict(cell=(124, 112), ctx=(232, 112), rows=8), LARGE: dict(cell=(250, 140), ctx=(250, 140), rows=5)}
KIND = {"img": "image", "bg": "CSS background", "pseudo": "CSS icon", "svg": "inline SVG", "svgimage": "SVG image"}


# ---------- colour maths (same as the runtime) ----------
def lin(u):
    u = u / 255.0
    return np.where(u <= 0.04045, u / 12.92, ((u + 0.055) / 1.055) ** 2.4)


def oklab(rgb):  # (..., 3) uint8 -> (..., 3)
    r, g, b = lin(rgb[..., 0]), lin(rgb[..., 1]), lin(rgb[..., 2])
    l = np.cbrt(0.4122214708 * r + 0.5363325363 * g + 0.0514459929 * b)
    m = np.cbrt(0.2119034982 * r + 0.6806995451 * g + 0.1073969566 * b)
    s = np.cbrt(0.0883024619 * r + 0.2817188376 * g + 0.6299787005 * b)
    return np.stack([0.2104542553 * l + 0.7936177850 * m - 0.0040720468 * s,
                     1.9779984951 * l - 2.4285922050 * m + 0.4505937099 * s,
                     0.0259040371 * l + 0.7827717662 * m - 0.8086757660 * s], -1)


def to_rgb(L, a, b):
    l_ = L + 0.3963377774 * a + 0.2158037573 * b
    m_ = L - 0.1055613458 * a - 0.0638541728 * b
    s_ = L - 0.0894841775 * a - 1.2914855480 * b
    l, m, s = l_ ** 3, m_ ** 3, s_ ** 3
    def enc(x):
        x = np.clip(x, 0, 1)
        return np.round(255 * np.where(x <= 0.0031308, 12.92 * x, 1.055 * np.power(x, 1 / 2.4) - 0.055))
    return np.stack([enc(4.0767416621 * l - 3.3077115913 * m + 0.2309699292 * s),
                     enc(-1.2684380046 * l + 2.6097574011 * m - 0.3413193965 * s),
                     enc(-0.0041960863 * l - 0.7034186147 * m + 1.7076147010 * s)], -1)


def invert(rgba):
    x = 1.0 - rgba[..., :3] / 255.0
    out = rgba.copy()
    out[..., :3] = np.round(np.clip(x @ HUE180.T, 0, 1) * 255)
    return out


def adapt(rgba):
    """Port of nadir-map.js adaptPixels (incl. transparent mode). None when there is nothing to adapt."""
    d = rgba.astype(np.float64)
    h, w = d.shape[:2]
    ring = np.zeros((h, w), bool)
    ring[0, :] = ring[-1, :] = ring[:, 0] = ring[:, -1] = True
    rp = d[ring]
    ringN = len(rp)
    opq = rp[:, 3] >= 250
    n = int(opq.sum())
    o = oklab(rp[opq][:, :3]) if n else np.zeros((0, 3))
    lightm = (o[:, 0] > 0.9) & (np.hypot(o[:, 1], o[:, 2]) < 0.04) if n else np.zeros(0, bool)
    light = int(lightm.sum())
    bg = None
    if n and light / n >= 0.5:
        bg = oklab(rp[opq][lightm][:, :3].mean(0))
    elif n > 0.2 * ringN:
        return None
    else:
        a = d[..., 3] >= 128
        if not a.any():
            return None
        ol = oklab(d[a][:, :3])
        ink = (ol[:, 0] < 0.5) & (np.hypot(ol[:, 1], ol[:, 2]) < 0.05)
        if ink.mean() < 0.08:
            return None
    out = np.zeros_like(d)
    vis = d[..., 3] > 0
    px = d[vis]
    ol = oklab(px[:, :3])
    if bg is not None:
        dist = np.linalg.norm(ol - bg, axis=1)
        cov = np.clip((dist - 0.03) / 0.22, 0, 1)
    else:
        cov = np.ones(len(px))
    keep = cov > 0
    c = np.where(keep, cov, 1.0)
    if bg is not None:
        L = (ol[:, 0] - (1 - c) * bg[0]) / c
        A = (ol[:, 1] - (1 - c) * bg[1]) / c
        B = (ol[:, 2] - (1 - c) * bg[2]) / c
    else:
        L, A, B = ol[:, 0].copy(), ol[:, 1].copy(), ol[:, 2].copy()
    L = np.clip(L, 0, 1)
    C = np.hypot(A, B)
    neutral = C < 0.05
    lift = (~neutral) & (L < 0.62)
    L = np.where(neutral, 0.94 - 0.75 * L, np.where(lift, 0.62 + 0.15 * L / 0.62, L))
    A = np.where(lift, A * 0.92, A)
    B = np.where(lift, B * 0.92, B)
    rgb = to_rgb(L, A, B)
    res = np.zeros((len(px), 4))
    res[:, :3] = rgb
    res[:, 3] = np.round(px[:, 3] * cov)
    res[~keep] = 0
    out[vis] = res
    return out.astype(np.uint8)


# ---------- page data ----------
def hex_rgba(h):
    h = h.lstrip("#")
    r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)
    a = int(h[6:8], 16) / 255 if len(h) >= 8 else 1.0
    return r, g, b, a


def over(top, base):
    r, g, b, a = top
    return tuple(round(a * t + (1 - a) * u) for t, u in zip((r, g, b), base))


def dark_palette(page, design):
    """Dark surface per token index and the main dark text colour, from the Haiku design."""
    cols = {int(k): v for k, v in design["colors"].items()} if isinstance(design["colors"], dict) else {int(c["id"].lstrip("t")): c["dark"] for c in design["colors"]}
    toks = page["tokens"]
    root = next((i for i, t in enumerate(toks) if t["role"] == 0 and t["f"][24] > 0.5), 0)
    canvas = hex_rgba(cols[root])[:3] if root in cols else (18, 18, 20)
    best, text = -1, (230, 232, 235)
    for i, t in enumerate(toks):
        if t["role"] == 1 and i in cols and t["f"][19] > best:
            best, text = t["f"][19], over(hex_rgba(cols[i]), canvas)
    surf = {i: over(hex_rgba(cols[i]), canvas) for i, t in enumerate(toks) if i in cols and t["role"] == 0}
    return surf, canvas, text


def content_area(o, big):
    """The big thumbnail is letterboxed into a square: cut out the image itself."""
    S = big.shape[0]
    if o.get("cropPx"):
        nw, nh = o["cropPx"][2], o["cropPx"][3]
    elif o.get("kind") == "svg":
        nw, nh = o["w"], o["h"]
    else:
        nw, nh = (o.get("nw") or o["w"]), (o.get("nh") or o["h"])
    nw, nh = max(1, nw), max(1, nh)
    s = min(S / nw, S / nh)
    dw, dh = max(1, round(nw * s)), max(1, round(nh * s))
    x0, y0 = (S - dw) // 2, (S - dh) // 2
    return big[y0:y0 + dh, x0:x0 + dw]


def decode_png(data_url):
    return np.array(Image.open(io.BytesIO(base64.b64decode(data_url.split(",", 1)[1]))).convert("RGBA"))


def bucket(o):
    s = o.get("stats") or {}
    area = o["w"] * o["h"]
    if o.get("logo") or (o.get("homeLink") and o.get("header")):
        return "logo"
    if o["kind"] == "bg" and area >= 20000:
        return "backdrop"
    if o["kind"] in ("pseudo",) or o.get("cropPx"):
        return "icon"
    if s.get("transp", 0) > 0.15:
        return "transparent"
    if s.get("transp", 1) < 0.02 and s.get("colors", 0) > 0.18 and s.get("stdL", 0) > 0.12:
        return "photo"
    return "flat"


RATE = {"logo": 1.0, "backdrop": 1.0, "icon": 1.0, "transparent": 1.0, "flat": 0.8, "photo": 0.12}


def candidates(task):
    """All sampled elements of one page with everything needed to draw their rows."""
    crawl, haiku, pid, tag, seed = task
    rnd = random.Random(seed)
    try:
        page = json.load(gzip.open(Path(crawl) / "pages" / pid / "light.json.gz"))
        design = json.load(open(Path(haiku) / f"{pid}.json"))
        if design.get("refusal") or not design.get("colors"):
            return []
        surf, canvas, text = dark_palette(page, design)
    except Exception:
        return []
    out = []
    for i, o in enumerate(page["images"]):
        if not o.get("big") or o["w"] * o["h"] < 64 or o["w"] < 4 or o["h"] < 4:
            continue
        if o["kind"] == "svg" and not o.get("useExt"):  # inline SVGs are recoloured by the colour model, not as images
            continue
        b = bucket(o)
        if rnd.random() > RATE[b]:
            continue
        dsurf = surf.get(o.get("underTok", -1), canvas)
        out.append({"page": pid, "crawl": tag, "img": i, "bucket": b, "o": {k: o.get(k) for k in ("kind", "src", "x", "y", "w", "h", "nw", "nh", "cropPx", "logo", "homeLink", "header", "inLink", "surf", "big", "stats")},
                    "dsurf": dsurf, "dtext": text, "light": tuple((o.get("surf") or [255, 255, 255])[:3])})
    return out


# ---------- drawing ----------
def font(sz):
    try:
        return ImageFont.load_default(size=sz)
    except TypeError:
        return ImageFont.load_default()


F11, F13, F15 = font(11), font(13), font(15)


def fit(o, box, maxzoom=4.0):
    s = min(box[0] / max(1, o["w"]), box[1] / max(1, o["h"]), maxzoom)
    return max(1, round(o["w"] * s)), max(1, round(o["h"] * s)), s


def variant_cell(img, o, typ, size, surface, text, mode):
    W, H = size
    cell = Image.new("RGBA", (W, H), tuple(surface) + (255,))
    if mode != "D":
        dw, dh, _ = fit(o, (W - 12, H - 12))
        im = Image.fromarray(img, "RGBA").resize((dw, dh), Image.LANCZOS)
        cell.alpha_composite(im, ((W - dw) // 2, (H - dh) // 2))
    if typ == LARGE and o["kind"] == "bg":  # content on a background image: in dark mode its text is light
        dr = ImageDraw.Draw(cell)
        dr.text((12, H // 2 - 22), "Heading on this background", fill=tuple(text), font=F15)
        dr.text((12, H // 2 + 2), "Body text that sits on top of it.", fill=tuple(text), font=F13)
    return cell


def context_cell(shot, e, size):
    W, H = size
    o = e["o"]
    if shot is not None and o["y"] + o["h"] <= shot.height and o["x"] >= -2 and o["x"] + o["w"] <= shot.width + 2:
        m = max(30, 0.75 * max(o["w"], o["h"]))
        x0, y0 = max(0, o["x"] - m), max(0, o["y"] - m)
        x1, y1 = min(shot.width, o["x"] + o["w"] + m), min(shot.height, o["y"] + o["h"] + m)
        crop = shot.crop((int(x0), int(y0), int(x1), int(y1)))
        s = min(W / crop.width, H / crop.height, 3.0)
        cw, ch = max(1, round(crop.width * s)), max(1, round(crop.height * s))
        crop = crop.resize((cw, ch), Image.LANCZOS).convert("RGBA")
        cell = Image.new("RGBA", (W, H), (236, 236, 236, 255))
        ox, oy = (W - cw) // 2, (H - ch) // 2
        cell.alpha_composite(crop, (ox, oy))
        dr = ImageDraw.Draw(cell)
        rx, ry = ox + (o["x"] - x0) * s, oy + (o["y"] - y0) * s
        dr.rectangle([rx - 2, ry - 2, rx + o["w"] * s + 1, ry + o["h"] * s + 1], outline=(230, 20, 20, 255), width=2)
        return cell
    # below the captured screenshot: the element on its light surface
    cell = Image.new("RGBA", (W, H), tuple(e["light"]) + (255,))
    big = content_area(o, decode_png(o["big"]))
    dw, dh, _ = fit(o, (W - 16, H - 16))
    cell.alpha_composite(Image.fromarray(big, "RGBA").resize((dw, dh), Image.LANCZOS), ((W - dw) // 2, (H - dh) // 2))
    ImageDraw.Draw(cell).text((4, H - 14), "(on its light background)", fill=(90, 90, 90), font=F11)
    return cell


def draw_sheet(rows, typ, shots, path):
    g = GEOM[typ]
    cw, ch = g["cell"]
    xw, xh = g["ctx"]
    pad, lab, cap = 6, 30, 16
    W = lab + xw + 4 * (cw + pad) + pad
    rh = max(ch, xh) + cap + 10
    H = 26 + rh * len(rows)
    sheet = Image.new("RGB", (W, H), (255, 255, 255))
    dr = ImageDraw.Draw(sheet)
    heads = ["light page (red box)", "A keep", "B invert", "C adapt", "D remove"]
    xs = [lab, lab + xw + pad] + [lab + xw + pad + k * (cw + pad) for k in range(1, 4)]
    for x, t in zip(xs, heads):
        dr.text((x + 4, 6), t, fill=(0, 0, 0), font=F13)
    meta = []
    for r, e in enumerate(rows):
        y = 26 + r * rh
        o = e["o"]
        dr.text((4, y + ch // 2 - 8), f"#{r + 1}", fill=(0, 0, 0), font=F15)
        sheet.paste(context_cell(shots.get(e["page"]), e, (xw, xh)).convert("RGB"), (xs[0], y))
        big = content_area(o, decode_png(o["big"]))
        offers = "AB"
        inv = invert(big)
        ad = adapt(big)
        cells = [variant_cell(big, o, typ, (cw, ch), e["dsurf"], e["dtext"], "A"), variant_cell(inv, o, typ, (cw, ch), e["dsurf"], e["dtext"], "B")]
        if ad is not None:
            offers += "C"
            cells.append(variant_cell(ad, o, typ, (cw, ch), e["dsurf"], e["dtext"], "C"))
        else:
            cells.append(None)
        if typ == LARGE and o["kind"] == "bg":
            offers += "D"
            cells.append(variant_cell(big, o, typ, (cw, ch), e["dsurf"], e["dtext"], "D"))
        else:
            cells.append(None)
        for k, c in enumerate(cells):
            x = xs[1 + k]
            if c is None:
                dr.rectangle([x, y, x + cw - 1, y + ch - 1], fill=(225, 225, 225))
                dr.text((x + cw // 2 - 14, y + ch // 2 - 7), "n/a", fill=(140, 140, 140), font=F13)
            else:
                sheet.paste(c.convert("RGB"), (x, y))
        flags = [f for f, v in (("logo", o.get("logo")), ("links home", o.get("homeLink")), ("in a link", o.get("inLink") and not o.get("homeLink")), ("in header", o.get("header"))) if v]
        _, _, z = fit(o, (cw - 12, ch - 12))
        zoom = f", shown {z:.1f}x" if abs(z - 1) > 0.15 else ""
        dr.text((xs[1], y + ch + 2), f"#{r + 1}: {KIND.get(o['kind'], o['kind'])} {o['w']}x{o['h']} px{zoom}" + (" · " + ", ".join(flags) if flags else ""), fill=(60, 60, 60), font=F11)
        meta.append({"k": r + 1, "page": e["page"], "crawl": e["crawl"], "img": e["img"], "bucket": e["bucket"], "offer": offers, "kind": o["kind"]})
    sheet.save(path, optimize=True)
    return meta


def render_chunk(job):
    wid, rows_by_sheet, crawls, out = job
    shots = {}
    outp = Path(out)
    lines = []
    for sid, typ, rows in rows_by_sheet:
        for e in rows:
            if e["page"] not in shots:
                p = Path(crawls[e["crawl"]]) / "pages" / e["page"] / "shot.jpg"
                try:
                    shots[e["page"]] = Image.open(p).convert("RGB")
                except Exception:
                    shots[e["page"]] = None
        try:
            meta = draw_sheet(rows, typ, shots, outp / "sheets" / f"{sid}.png")
            lines.append(json.dumps({"sheet": sid, "type": typ, "rows": meta}))
        except Exception as ex:
            print("sheet error", sid, repr(ex)[:200], flush=True)
        if len(shots) > 64:
            shots.clear()
    with open(outp / f"meta.{wid}.jsonl", "a") as f:
        f.write("\n".join(lines) + ("\n" if lines else ""))
    return len(lines)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--crawl", action="append", required=True, help="TAG=CRAWL_DIR=HAIKU_DIR (repeatable)")
    ap.add_argument("--out", required=True)
    ap.add_argument("--max-elements", type=int, default=150000)
    ap.add_argument("--max-per-src", type=int, default=2)
    ap.add_argument("--workers", type=int, default=28)
    ap.add_argument("--exclude-pages", default="", help="file of page ids already done")
    ap.add_argument("--flat-rate", type=float, default=None, help="sampling rate for opaque flat/photo-like images (mostly 'keep')")
    args = ap.parse_args()
    if args.flat_rate is not None:
        RATE["flat"] = args.flat_rate
    out = Path(args.out)
    (out / "sheets").mkdir(parents=True, exist_ok=True)
    crawls, tasks = {}, []
    done = set(l.strip() for l in open(args.exclude_pages)) if args.exclude_pages and os.path.exists(args.exclude_pages) else set()
    for spec in args.crawl:
        tag, cdir, hdir = spec.split("=")
        crawls[tag] = cdir
        for hp in Path(hdir).glob("*.json"):
            pid = hp.stem
            if len(pid) == 16 and pid not in done and (Path(cdir) / "pages" / pid / "light.json.gz").exists():
                tasks.append((cdir, hdir, pid, tag, int(hashlib.md5(pid.encode()).hexdigest()[:8], 16)))
    random.Random(3).shuffle(tasks)
    print(len(tasks), "pages with designs", flush=True)
    per_src, elems, by_bucket = {}, [], {}
    with Pool(args.workers) as pool:
        for res in pool.imap_unordered(candidates, tasks, chunksize=16):
            for e in res:
                s = e["o"]["src"]
                if per_src.get(s, 0) >= args.max_per_src:
                    continue
                per_src[s] = per_src.get(s, 0) + 1
                elems.append(e)
                by_bucket[e["bucket"]] = by_bucket.get(e["bucket"], 0) + 1
            if len(elems) >= args.max_elements:
                break
    elems = elems[: args.max_elements]
    print(len(elems), "elements", by_bucket, flush=True)
    with open(out / "pages_used.txt", "a") as f:
        f.write("\n".join(sorted(set(e["page"] for e in elems))) + "\n")
    # sheets: small and large elements apart; rows mixed across pages
    random.Random(5).shuffle(elems)
    groups = {SMALL: [], LARGE: []}
    for e in elems:
        o = e["o"]
        typ = LARGE if (o["w"] * o["h"] >= 20000 or o["w"] > 300) else SMALL
        groups[typ].append(e)
    sheets = []
    for typ, es in groups.items():
        n = GEOM[typ]["rows"]
        for i in range(0, len(es), n):
            rows = es[i:i + n]
            sid = hashlib.sha1(json.dumps([(r["page"], r["img"]) for r in rows]).encode()).hexdigest()[:16]
            sheets.append((sid, typ, rows))
    jobs = [(w, sheets[w::args.workers], crawls, str(out)) for w in range(args.workers)]
    with Pool(args.workers) as pool:
        n = sum(pool.map(render_chunk, jobs))
    print("sheets", n, flush=True)


if __name__ == "__main__":
    main()