#!/usr/bin/env python3
"""Repaint a GLB's baked vertex colour (`COLOR_0`) into a named palette — zero draws, zero triangles.

WHY THIS EXISTS (R42 §42.1 + §42.2, "the pub props render as untextured white blobs")
--------------------------------------------------------------------------------------
R41's own note said "37 of 39 have no texture at all". Measuring the shipped files says something
different and more useful: **all 37 DO carry COLOR_0 and it does render** — it is just the wrong
colour. Mean vertex colour across the set is a bleached near-grey (15 of 37 sit below 0.05
saturation; `beer_umbrella` is sRGB 238,238,238 and `jukebox` is 175,174,173), because the source
library is photogrammetry/generative output whose "colour" is really baked ambient occlusion. The
fault is therefore not a missing map — it is a missing palette.

The cheapest correct fix is consequently NOT a texture. Lane C took the interior from 14,140 draws
to 117 by collapsing multi-material primitives into one `COLOR_0` primitive each (`merge_prims.py`);
adding per-asset baseColour maps would re-import the exact cost that bought. So this tool rewrites
the COLOR_0 buffer *in place*, preserving the accessor's component type and element count — same
byte length, same primitives, same triangles, **same draw count** — and only touches the material
block to set metallic/roughness/emissive (material count does not affect draw count; three.js issues
one draw per glTF *primitive*, and `batch.js` never descends into a GLB instance because the loaded
root is a Group, not a Mesh).

Four routes, chosen per asset by what the source colour actually is:

  ramp    luminance -> a 2..5 stop palette gradient. For the bleached-grey majority: the greyscale
          IS the baked shading, so mapping it through a dark->light palette keeps every AO gradient
          and adds hue. Normalisation is per-asset between the source's own 1st/99th luminance
          percentiles, so a flat-grey asset still spans the whole ramp.
  lut     exact source-colour -> target, ordered rules, first match wins, optional position
          predicates (`cyl`, `y_min`, `y_max`). For the DJ gear, whose colours are a discrete
          5..12 entry palette where each entry IS a part (chassis / platter / slipmat / arm).
  tint    keep the source hue relationship, force saturation and remap value. For the minority that
          already carry meaningful colour (`pokie`'s lit artwork, `planter`'s foliage).
  huesplit  two ramps selected by source hue — foliage vs pot, for the planted props.
  factor  no COLOR_0 at all (a real texture): tint via `baseColorFactor`, which multiplies the map.

    python3 pipeline/palette.py --spec pipeline/r42_palette.json --src DIR --out DIR
    python3 pipeline/palette.py --spec pipeline/r42_palette.json --src DIR --out DIR --only jukebox
    python3 pipeline/palette.py --report DIR/*.glb        # mean sRGB / saturation / prims / draws

RE-RUN SAFETY: `ramp` and `tint` normalise against the source file's own luminance spread, so a
second pass over an already-painted file is NOT a no-op. Every output is stamped
`asset.extras.procity_palette`, and a stamped input is skipped unless `--force`. The pristine
pre-R42 GLBs are kept at `pipeline/.props_orig/r42/` (git-ignored, like `_normalized/` itself),
which is the directory to point `--src` at when re-running.
"""
import json, os, struct, sys, colorsys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from merge_prims import read_glb, write_glb, COMPONENT, NCOMP     # noqa: E402


# ── colour helpers ───────────────────────────────────────────────────────────────────────────────
def srgb_to_lin(c):
    return c / 12.92 if c <= 0.04045 else ((c + 0.055) / 1.055) ** 2.4


def lin_to_srgb(c):
    c = max(0.0, min(1.0, c))
    return 12.92 * c if c <= 0.0031308 else 1.055 * (c ** (1 / 2.4)) - 0.055


def hex_lin(h):
    h = h.lstrip("#")
    return tuple(srgb_to_lin(int(h[i:i + 2], 16) / 255.0) for i in (0, 2, 4))


def lum(rgb):
    return 0.2126 * rgb[0] + 0.7152 * rgb[1] + 0.0722 * rgb[2]


def lerp3(a, b, t):
    return tuple(a[k] + (b[k] - a[k]) * t for k in range(3))


# ── glb attribute io ─────────────────────────────────────────────────────────────────────────────
def acc_layout(gltf, ai):
    a = gltf["accessors"][ai]
    fmt, sz = COMPONENT[a["componentType"]]
    n = NCOMP[a["type"]]
    bv = gltf["bufferViews"][a["bufferView"]]
    base = bv.get("byteOffset", 0) + a.get("byteOffset", 0)
    stride = bv.get("byteStride") or sz * n
    return a, fmt, sz, n, base, stride


def read_vec(gltf, binc, ai):
    a, fmt, sz, n, base, stride = acc_layout(gltf, ai)
    return [struct.unpack_from("<" + fmt * n, binc, base + i * stride) for i in range(a["count"])]


def write_colors(gltf, buf, ai, rows):
    """Write linear RGB back into an existing COLOR_0 accessor. Byte length never changes."""
    a, fmt, sz, n, base, stride = acc_layout(gltf, ai)
    ct = a["componentType"]
    div = 65535.0 if ct == 5123 else (255.0 if ct == 5121 else 1.0)
    for i, rgb in enumerate(rows):
        old = struct.unpack_from("<" + fmt * n, buf, base + i * stride)
        if div == 1.0:
            vals = [float(c) for c in rgb]
        else:
            vals = [int(round(max(0.0, min(1.0, c)) * div)) for c in rgb]
        if n == 4:
            vals.append(old[3])
        struct.pack_into("<" + fmt * n, buf, base + i * stride, *vals)


# ── the four routes ──────────────────────────────────────────────────────────────────────────────
def route_ramp(cols, spec):
    stops = [(float(t), hex_lin(h)) for t, h in spec["stops"]]
    stops.sort()
    L = sorted(lum(c) for c in cols)
    lo = L[max(0, int(len(L) * 0.01))]
    hi = L[min(len(L) - 1, int(len(L) * 0.99))]
    span = max(hi - lo, 1e-4)
    gamma = float(spec.get("gamma", 1.0))
    out = []
    for c in cols:
        t = min(1.0, max(0.0, (lum(c) - lo) / span)) ** gamma
        for k in range(len(stops) - 1):
            t0, c0 = stops[k]
            t1, c1 = stops[k + 1]
            if t <= t1 or k == len(stops) - 2:
                f = 0.0 if t1 <= t0 else min(1.0, max(0.0, (t - t0) / (t1 - t0)))
                out.append(lerp3(c0, c1, f))
                break
        else:
            out.append(stops[-1][1])
    return out


def route_tint(cols, spec):
    hue = spec.get("hue")                       # degrees, or None to keep source hue
    hue_rot = float(spec.get("hue_rot", 0.0)) / 360.0
    sat_mul = float(spec.get("sat_mul", 1.0))
    sat_set = spec.get("sat")
    v0, v1 = spec.get("val", [0.0, 1.0])
    S = sorted(max(lin_to_srgb(x) for x in c) for c in cols)
    vlo = S[max(0, int(len(S) * 0.01))]
    vhi = S[min(len(S) - 1, int(len(S) * 0.99))]
    vspan = max(vhi - vlo, 1e-4)
    out = []
    for c in cols:
        r, g, b = (lin_to_srgb(x) for x in c)
        h, s, v = colorsys.rgb_to_hsv(r, g, b)
        if hue is not None:
            h = (float(hue) / 360.0)
        h = (h + hue_rot) % 1.0
        s = float(sat_set) if sat_set is not None else min(1.0, s * sat_mul)
        v = v0 + (min(1.0, max(0.0, (v - vlo) / vspan)) * (v1 - v0))
        r, g, b = colorsys.hsv_to_rgb(h, s, v)
        out.append((srgb_to_lin(r), srgb_to_lin(g), srgb_to_lin(b)))
    return out


def route_huesplit(cols, spec):
    """Two ramps chosen by SOURCE hue — 'a planter is a pot and a plant, and they are not one ramp'."""
    lo_h, hi_h = spec.get("hue_band", [60, 190])
    idx_a, idx_b = [], []
    for i, c in enumerate(cols):
        r, g, b = (lin_to_srgb(x) for x in c)
        h, s, v = colorsys.rgb_to_hsv(r, g, b)
        deg = h * 360.0
        (idx_a if (lo_h <= deg <= hi_h and s >= spec.get("hue_min_sat", 0.05)) else idx_b).append(i)
    out = [None] * len(cols)
    for idxs, key in ((idx_a, "in_stops"), (idx_b, "out_stops")):
        if not idxs:
            continue
        sub = route_ramp([cols[i] for i in idxs], {"stops": spec[key], "gamma": spec.get("gamma", 1.0)})
        for j, i in enumerate(idxs):
            out[i] = sub[j]
    return [o if o is not None else c for o, c in zip(out, cols)]


def route_lut(cols, pos, spec):
    rules = spec["rules"]
    default = hex_lin(spec["default"]) if spec.get("default") else None
    out = []
    for c, p in zip(cols, pos):
        hit = None
        for rl in rules:
            if "src" in rl and max(abs(c[k] - rl["src"][k]) for k in range(3)) > 2e-3:
                continue
            if "y_min" in rl and p[1] < rl["y_min"]:
                continue
            if "y_max" in rl and p[1] > rl["y_max"]:
                continue
            cyl = rl.get("cyl")
            if cyl:
                cx, cz = cyl["center"]
                r = ((p[0] - cx) ** 2 + (p[2] - cz) ** 2) ** 0.5
                if r > cyl.get("r_max", 1e9) or r < cyl.get("r_min", -1.0):
                    continue
            hit = rl
            break
        out.append(hex_lin(hit["to"]) if hit else (default if default else c))
    return out


ROUTES = {"ramp", "lut", "tint", "huesplit", "factor", "none"}


# ── per-asset application ────────────────────────────────────────────────────────────────────────
def mat_for(gltf, prim, want):
    """Return a material index carrying `want` (metallic/roughness/emissive), cloning if shared."""
    mi = prim.get("material")
    if mi is None:
        return None
    mats = gltf["materials"]
    cur = mats[mi]
    users = sum(1 for m in gltf.get("meshes", []) for p in m["primitives"] if p.get("material") == mi)
    tgt = json.loads(json.dumps(cur))
    pbr = tgt.setdefault("pbrMetallicRoughness", {})
    if "metallic" in want:
        pbr["metallicFactor"] = float(want["metallic"])
    if "roughness" in want:
        pbr["roughnessFactor"] = float(want["roughness"])
    if "factor" in want:
        pbr["baseColorFactor"] = list(hex_lin(want["factor"])) + [1.0]
    if "emissive" in want:
        tgt["emissiveFactor"] = list(hex_lin(want["emissive"]))
    if tgt == cur:
        return mi
    if users == 1:
        mats[mi] = tgt
        return mi
    mats.append(tgt)
    return len(mats) - 1


STAMP = "procity_palette"


def stamped(gltf):
    """`ramp` and `tint` normalise against the SOURCE's own luminance spread, so running them twice
    is not a no-op — the second pass re-stretches an already-stretched file. The asset therefore
    carries a stamp and a second run refuses unless --force. (`lut`/`factor` are idempotent; the
    stamp is applied uniformly so one rule covers the whole spec.)"""
    return (gltf.get("asset", {}).get("extras") or {}).get(STAMP)


def apply_asset(path, out_path, spec, force=False):
    gltf, binc = read_glb(path)
    was = stamped(gltf)
    if was and not force:
        return -1
    buf = bytearray(binc)
    touched = 0
    seen_acc = set()
    for mesh in gltf.get("meshes", []):
        mname = mesh.get("name", "")
        for prim in mesh["primitives"]:
            sub = dict(spec)
            sub.pop("meshes", None)
            over = (spec.get("meshes") or {}).get(mname)
            if over:
                sub.update(over)
            route = sub.get("route", "none")
            if route not in ROUTES:
                raise ValueError(f"{path}: unknown route {route!r}")
            want = {k: sub[k] for k in ("metallic", "roughness", "emissive", "factor") if k in sub}
            if want:
                prim["material"] = mat_for(gltf, prim, want)
            ai = prim["attributes"].get("COLOR_0")
            if route in ("none", "factor") or ai is None:
                continue
            if ai in seen_acc:
                print(f"  ! {os.path.basename(path)}: COLOR_0 accessor {ai} shared across primitives; "
                      f"already written, skipping second pass")
                continue
            seen_acc.add(ai)
            cols = [tuple(v[:3]) for v in read_vec(gltf, binc, ai)]
            if route == "ramp":
                new = route_ramp(cols, sub)
            elif route == "tint":
                new = route_tint(cols, sub)
            elif route == "huesplit":
                new = route_huesplit(cols, sub)
            else:
                new = route_lut(cols, read_vec(gltf, binc, prim["attributes"]["POSITION"]), sub)
            write_colors(gltf, buf, ai, new)
            touched += 1
    gltf.setdefault("asset", {}).setdefault("extras", {})[STAMP] = spec.get("stamp", "r42")
    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
    write_glb(out_path, gltf, bytes(buf))
    return touched


# ── reporting ────────────────────────────────────────────────────────────────────────────────────
def report(paths):
    print(f"{'file':44} {'prims':>5} {'mats':>4} {'tris':>6} {'meanRGB':16} {'lum':>5} {'sat':>5}")
    for p in paths:
        gltf, binc = read_glb(p)
        prims = tris = 0
        tot = [0.0, 0.0, 0.0]
        n = 0
        for mesh in gltf.get("meshes", []):
            for prim in mesh["primitives"]:
                prims += 1
                if "indices" in prim:
                    tris += gltf["accessors"][prim["indices"]]["count"] // 3
                ai = prim["attributes"].get("COLOR_0")
                if ai is None:
                    continue
                a, fmt, sz, nc, base, stride = acc_layout(gltf, ai)
                div = 65535.0 if a["componentType"] == 5123 else (255.0 if a["componentType"] == 5121 else 1.0)
                for v in read_vec(gltf, binc, ai):
                    for k in range(3):
                        tot[k] += v[k] / div
                    n += 1
        if n:
            mean = [t / n for t in tot]
            rgb = [round(lin_to_srgb(c) * 255) for c in mean]
            sat = (max(mean) - min(mean)) / max(mean) if max(mean) > 0 else 0.0
            s = f"{str(rgb):16} {lum(mean):5.3f} {sat:5.3f}"
        else:
            s = f"{'(textured)':16} {'-':>5} {'-':>5}"
        print(f"{os.path.basename(p):44} {prims:5d} {len(gltf.get('materials', [])):4d} {tris:6d} {s}")


def main():
    av = sys.argv[1:]
    if "--report" in av:
        report([a for a in av[av.index("--report") + 1:] if a.endswith(".glb")])
        return 0
    spec_path = av[av.index("--spec") + 1]
    src = av[av.index("--src") + 1]
    out = av[av.index("--out") + 1]
    only = av[av.index("--only") + 1].split(",") if "--only" in av else None
    force = "--force" in av
    doc = json.load(open(spec_path))
    n = skipped = 0
    for aid, spec in doc["assets"].items():
        if only and aid not in only:
            continue
        f = spec["file"]
        touched = apply_asset(os.path.join(src, f), os.path.join(out, f), spec, force)
        if touched < 0:
            print(f"  {aid:18} SKIP — already stamped {STAMP}; pass --force to repaint a repaint")
            skipped += 1
            continue
        print(f"  {aid:18} {spec.get('route', 'none'):9} prims_recoloured={touched}  -> {f}")
        n += 1
    print(f"palette: {n} assets repainted, {skipped} skipped")
    return 0


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