#!/usr/bin/env python3
"""Collapse a multi-material GLB to ONE primitive per mesh by baking baseColorFactor -> COLOR_0.

WHY THIS EXISTS (R41 §41.2, the 3dstore import)
-----------------------------------------------
`~/Documents/3dstore/` was handed over as "45 of 46 pass the house law as-is". On tris, scale,
Draco and textures that is true. On the number that actually decides whether Lane C can place them
it is not: **14 of the files carry 16 materials on a single mesh**, one flat colour per record
sleeve (`Mat_Sleeve_<hex>`). glTF gives each material its own primitive, and three.js
`GLTFLoader` instantiates one `Mesh` per primitive — so a 1,020-tri "Black Tub" costs **16 draw
calls**, not one. The interior draw margin is 162. Ten tubs would be the entire budget.

The materials are pure `baseColorFactor` — no textures, no UVs, nothing but a flat colour and a
metallic/roughness pair. That is exactly what a vertex colour reproduces. So we merge every
primitive of a mesh into one, write each source primitive's `baseColorFactor` into a `COLOR_0`
attribute on its own vertices, and emit a single white material with the dominant primitive's
PBR values. glTF multiplies `COLOR_0` by `baseColorFactor`, and both live in LINEAR space, so a
white factor times the original linear colour is **numerically the original colour**. Byte-exact
look, 16 draws -> 1.

This is deliberately NOT `gltf-transform optimize` (ruling 5) and NOT a Blender re-export: the
geometry is already correct and every extra tool is another chance to move it. We rewrite the
container and touch nothing but the material binding.

    python3 pipeline/merge_prims.py IN.glb OUT.glb
    python3 pipeline/merge_prims.py --check OUT.glb          # report prims/draws/materials
"""
import json, os, struct, sys

COMPONENT = {5120: ("b", 1), 5121: ("B", 1), 5122: ("h", 2),
             5123: ("H", 2), 5125: ("I", 4), 5126: ("f", 4)}
NCOMP = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4, "MAT4": 16}


# ── container io ─────────────────────────────────────────────────────────────────────────────────
def read_glb(path):
    d = open(path, "rb").read()
    if d[:4] != b"glTF":
        raise ValueError(f"{path}: not a GLB")
    off, gltf, binc = 12, None, b""
    while off + 8 <= len(d):
        clen, ctype = struct.unpack("<II", d[off:off + 8])
        chunk = d[off + 8: off + 8 + clen]
        if ctype == 0x4E4F534A:
            gltf = json.loads(chunk.decode("utf-8"))
        elif ctype == 0x004E4942:
            binc = chunk
        off += 8 + clen
    return gltf, binc


def write_glb(path, gltf, binc):
    js = json.dumps(gltf, separators=(",", ":")).encode("utf-8")
    js += b" " * ((4 - len(js) % 4) % 4)
    binc += b"\x00" * ((4 - len(binc) % 4) % 4)
    total = 12 + 8 + len(js) + (8 + len(binc) if binc else 0)
    with open(path, "wb") as f:
        f.write(b"glTF" + struct.pack("<II", 2, total))
        f.write(struct.pack("<II", len(js), 0x4E4F534A) + js)
        if binc:
            f.write(struct.pack("<II", len(binc), 0x004E4942) + binc)


def read_accessor(gltf, binc, idx):
    """Return a list of tuples (or scalars) for accessor `idx`, honouring byteStride."""
    a = gltf["accessors"][idx]
    fmt, size = COMPONENT[a["componentType"]]
    n = NCOMP[a["type"]]
    count = a["count"]
    if "bufferView" not in a:                                  # all-zero accessor (spec-legal)
        return [(0,) * n for _ in range(count)] if n > 1 else [0] * count
    bv = gltf["bufferViews"][a["bufferView"]]
    base = bv.get("byteOffset", 0) + a.get("byteOffset", 0)
    stride = bv.get("byteStride") or (size * n)
    out = []
    for i in range(count):
        off = base + i * stride
        v = struct.unpack_from("<" + fmt * n, binc, off)
        out.append(v if n > 1 else v[0])
    return out


def norm_component(v, ctype):
    """glTF normalized-integer -> float, per the spec's exact divisors."""
    if ctype == 5121:
        return v / 255.0
    if ctype == 5123:
        return v / 65535.0
    if ctype == 5120:
        return max(v / 127.0, -1.0)
    if ctype == 5122:
        return max(v / 32767.0, -1.0)
    return float(v)


# ── the merge ────────────────────────────────────────────────────────────────────────────────────
def merge(gltf, binc):
    """Rewrite every multi-primitive mesh as one primitive with baked COLOR_0. Returns a report."""
    mats = gltf.get("materials", [])
    new_bin = bytearray()
    new_views, new_accs = [], []

    def add_accessor(data, fmt, size, ncomp, atype, ctype, minmax=False, target=None):
        while len(new_bin) % 4:
            new_bin.append(0)
        off = len(new_bin)
        for v in data:
            new_bin.extend(struct.pack("<" + fmt * ncomp, *(v if ncomp > 1 else (v,))))
        bv = {"buffer": 0, "byteOffset": off, "byteLength": len(new_bin) - off}
        if target:
            bv["target"] = target
        new_views.append(bv)
        acc = {"bufferView": len(new_views) - 1, "componentType": ctype,
               "count": len(data), "type": atype}
        if minmax and data:
            cols = list(zip(*data)) if ncomp > 1 else [data]
            acc["min"] = [min(c) for c in cols]
            acc["max"] = [max(c) for c in cols]
        new_accs.append(acc)
        return len(new_accs) - 1

    report = []
    for mi, mesh in enumerate(gltf.get("meshes", [])):
        prims = mesh.get("primitives", [])
        if not prims:
            continue
        pos, nrm, col, idx = [], [], [], []
        tri_by_prim = []
        for p in prims:
            attrs = p["attributes"]
            P = read_accessor(gltf, binc, attrs["POSITION"])
            N = (read_accessor(gltf, binc, attrs["NORMAL"])
                 if "NORMAL" in attrs else [(0.0, 1.0, 0.0)] * len(P))
            # existing COLOR_0 multiplies the factor (spec); normalize ints first
            if "COLOR_0" in attrs:
                ca = gltf["accessors"][attrs["COLOR_0"]]
                raw = read_accessor(gltf, binc, attrs["COLOR_0"])
                ct, nc = ca["componentType"], NCOMP[ca["type"]]
                existing = [tuple(norm_component(x, ct) for x in (v if nc > 1 else (v,)))
                            for v in raw]
                existing = [(e + (1.0,)) if len(e) == 3 else e for e in existing]
            else:
                existing = None
            m = mats[p["material"]] if p.get("material") is not None and mats else {}
            bcf = (m.get("pbrMetallicRoughness", {}) or {}).get("baseColorFactor", [1, 1, 1, 1])
            bcf = tuple(float(x) for x in (list(bcf) + [1, 1, 1, 1])[:4])
            base = len(pos)
            pos.extend(P)
            nrm.extend(N)
            if existing:
                col.extend(tuple(bcf[k] * existing[i][k] for k in range(4))
                           for i in range(len(P)))
            else:
                col.extend([bcf] * len(P))
            if "indices" in p:
                I = read_accessor(gltf, binc, p["indices"])
            else:
                I = list(range(len(P)))
            idx.extend(i + base for i in I)
            tri_by_prim.append((len(I) // 3, p.get("material")))

        # PBR donor = the material covering the most triangles, so the dominant surface keeps its
        # exact metallic/roughness. (On the 3dstore tubs that is a sleeve: metallic 0, rough 0.6.)
        donor_mat = max(tri_by_prim, key=lambda t: t[0])[1] if tri_by_prim else None
        donor = mats[donor_mat] if donor_mat is not None and mats else {}
        dp = donor.get("pbrMetallicRoughness", {}) or {}

        ctype = 5125 if len(pos) > 65535 else 5123
        fmt = "I" if ctype == 5125 else "H"
        i_acc = add_accessor(idx, fmt, 4 if ctype == 5125 else 2, 1, "SCALAR", ctype,
                             target=34963)
        p_acc = add_accessor(pos, "f", 4, 3, "VEC3", 5126, minmax=True, target=34962)
        n_acc = add_accessor(nrm, "f", 4, 3, "VEC3", 5126, target=34962)
        c_acc = add_accessor(col, "f", 4, 4, "VEC4", 5126, target=34962)

        mesh["primitives"] = [{
            "attributes": {"POSITION": p_acc, "NORMAL": n_acc, "COLOR_0": c_acc},
            "indices": i_acc, "material": 0, "mode": 4,
        }]
        report.append({"mesh": mi, "name": mesh.get("name"), "prims_before": len(prims),
                       "prims_after": 1, "tris": len(idx) // 3, "verts": len(pos)})

    gltf["materials"] = [{
        "name": "Merged_VertexColor",
        "doubleSided": any(m.get("doubleSided") for m in mats) if mats else False,
        "pbrMetallicRoughness": {
            "baseColorFactor": [1, 1, 1, 1],                     # identity: COLOR_0 carries it all
            "metallicFactor": dp.get("metallicFactor", 1.0),
            "roughnessFactor": dp.get("roughnessFactor", 1.0),
        },
    }]
    gltf["bufferViews"] = new_views
    gltf["accessors"] = new_accs
    gltf["buffers"] = [{"byteLength": len(new_bin)}]
    for k in ("images", "textures", "samplers"):                 # nothing references them now
        gltf.pop(k, None)
    gltf["extensionsUsed"] = [e for e in gltf.get("extensionsUsed", [])
                              if "draco" not in e.lower()] or None
    if not gltf["extensionsUsed"]:
        gltf.pop("extensionsUsed")
    return report, bytes(new_bin)


def check(path):
    gltf, _ = read_glb(path)
    meshes = gltf.get("meshes", [])
    nodes = gltf.get("nodes", [])
    draws = sum(len(meshes[n["mesh"]].get("primitives", [])) for n in nodes if "mesh" in n)
    prims = sum(len(m.get("primitives", [])) for m in meshes)
    tris = 0
    for m in meshes:
        for p in m.get("primitives", []):
            a = gltf["accessors"][p["indices"]] if "indices" in p else \
                gltf["accessors"][p["attributes"]["POSITION"]]
            tris += a["count"] // 3
    return {"file": os.path.basename(path), "draws": draws, "prims": prims,
            "materials": len(gltf.get("materials", [])), "tris": tris,
            "images": len(gltf.get("images", []))}


def main():
    args = sys.argv[1:]
    if args and args[0] == "--check":
        for f in args[1:]:
            print(check(f))
        return 0
    if len(args) != 2:
        print(__doc__)
        return 2
    src, dst = args
    gltf, binc = read_glb(src)
    before = check(src)
    rep, nb = merge(gltf, binc)
    os.makedirs(os.path.dirname(os.path.abspath(dst)), exist_ok=True)
    write_glb(dst, gltf, nb)
    after = check(dst)
    print(f"{os.path.basename(src):42} draws {before['draws']:>3} -> {after['draws']:>2} "
          f"| mats {before['materials']:>2} -> {after['materials']} "
          f"| tris {before['tris']} -> {after['tris']}"
          f"{'  *** TRI DRIFT ***' if before['tris'] != after['tris'] else ''}")
    return 0


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