#!/usr/bin/env python3
"""Dependency-free GLB inspector for the PROCITY asset audit.

Parses .glb binary containers (no Blender, no pygltflib) and reports, per file:
tri count, bounding box in metres (scale sanity), node/mesh/material/image counts,
and embedded texture pixel dimensions. Also flags likely house-GLB-law violations.

Usage:
    python3 glb_stat.py FILE.glb [FILE2.glb ...]          # human table
    python3 glb_stat.py --json DIR                        # JSON, recurse *.glb in DIR
    python3 glb_stat.py --json FILE.glb                   # JSON for one file

House GLB law (CITY_SPEC / LANE_E): metres, +Y up, origin at base, facing -Z,
props <=5k tris, textures <=1024, no Draco.
"""
import json, struct, sys, os, glob

COMPONENT_SIZE = {5120: 1, 5121: 1, 5122: 2, 5123: 2, 5125: 4, 5126: 4}
TYPE_COUNT = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4,
              "MAT2": 4, "MAT3": 9, "MAT4": 16}


def _read_glb(path):
    with open(path, "rb") as f:
        data = f.read()
    if data[:4] != b"glTF":
        # maybe a .gltf JSON, or non-glb; try JSON parse
        try:
            return json.loads(data.decode("utf-8")), b""
        except Exception:
            raise ValueError("not a GLB/glTF")
    ver, length = struct.unpack("<II", data[4:12])
    off, gltf, bin_chunk = 12, None, b""
    while off < len(data):
        clen, ctype = struct.unpack("<II", data[off:off + 8])
        chunk = data[off + 8: off + 8 + clen]
        if ctype == 0x4E4F534A:      # 'JSON'
            gltf = json.loads(chunk.decode("utf-8"))
        elif ctype == 0x004E4942:    # 'BIN\0'
            bin_chunk = chunk
        off += 8 + clen
    return gltf, bin_chunk


def _prim_tris(gltf, prim):
    """Triangle count for one primitive (mode 4 / default = TRIANGLES)."""
    mode = prim.get("mode", 4)
    accessors = gltf.get("accessors", [])
    if "indices" in prim:
        count = accessors[prim["indices"]]["count"]
    else:
        pos = prim.get("attributes", {}).get("POSITION")
        if pos is None:
            return 0
        count = accessors[pos]["count"]
    if mode == 4:
        return count // 3
    if mode in (5, 6):               # triangle strip / fan
        return max(0, count - 2)
    return 0                          # points/lines contribute no tris


def _mat_mul(A, B):
    return [[sum(A[i][k] * B[k][j] for k in range(4)) for j in range(4)] for i in range(4)]


def _node_matrix(n):
    if "matrix" in n:                       # column-major → row-major
        m = n["matrix"]
        return [[m[c * 4 + r] for c in range(4)] for r in range(4)]
    t = n.get("translation", [0, 0, 0]); q = n.get("rotation", [0, 0, 0, 1]); s = n.get("scale", [1, 1, 1])
    x, y, z, w = q
    R = [[1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w), 0],
         [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w), 0],
         [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y), 0],
         [0, 0, 0, 1]]
    S = [[s[0], 0, 0, 0], [0, s[1], 0, 0], [0, 0, s[2], 0], [0, 0, 0, 1]]
    T = [[1, 0, 0, t[0]], [0, 1, 0, t[1]], [0, 0, 1, t[2]], [0, 0, 0, 1]]
    return _mat_mul(_mat_mul(T, R), S)


def _xform(M, p):
    return [sum(M[i][j] * ([p[0], p[1], p[2], 1][j]) for j in range(4)) for i in range(3)]


def stat(path):
    gltf, _ = _read_glb(path)
    accessors = gltf.get("accessors", [])
    meshes = gltf.get("meshes", [])
    nodes = gltf.get("nodes", [])
    tris = 0
    for m in meshes:                        # tri count is transform-independent
        for p in m.get("primitives", []):
            tris += _prim_tris(gltf, p)
    # WORLD-space bbox: walk the scene graph so node TRS (unit scale / pivot offset) is respected —
    # local accessor min/max alone lies for multi-node assets (missed a 1km off-origin bug once).
    mn = [float("inf")] * 3
    mx = [float("-inf")] * 3
    scenes = gltf.get("scenes") or [{"nodes": list(range(len(nodes)))}]
    roots = scenes[gltf.get("scene", 0)].get("nodes", [])

    def walk(ni, M):
        n = nodes[ni]
        W = _mat_mul(M, _node_matrix(n))
        if "mesh" in n:
            for p in meshes[n["mesh"]].get("primitives", []):
                pos = p.get("attributes", {}).get("POSITION")
                if pos is None:
                    continue
                a = accessors[pos]
                if "min" not in a or "max" not in a:
                    continue
                lo, hi = a["min"], a["max"]
                for cx in (lo[0], hi[0]):
                    for cy in (lo[1], hi[1]):
                        for cz in (lo[2], hi[2]):
                            w = _xform(W, [cx, cy, cz])
                            for k in range(3):
                                mn[k] = min(mn[k], w[k]); mx[k] = max(mx[k], w[k])
        for c in n.get("children", []):
            walk(c, W)

    I = [[1 if i == j else 0 for j in range(4)] for i in range(4)]
    for r in roots:
        walk(r, I)
    dims = [round(mx[i] - mn[i], 3) if mx[i] > mn[i] else 0.0 for i in range(3)]
    # texture pixel sizes from embedded images (PNG/JPEG header sniff)
    tex_sizes = []
    images = gltf.get("images", [])
    bufviews = gltf.get("bufferViews", [])
    _, binc = _read_glb(path)
    for im in images:
        wh = None
        if "bufferView" in im and binc:
            bv = bufviews[im["bufferView"]]
            start = bv.get("byteOffset", 0)
            blob = binc[start:start + bv.get("byteLength", 0)]
            wh = _img_size(blob)
        if wh:
            tex_sizes.append(wh)
    exts = gltf.get("extensionsUsed", []) or []
    draco = any("draco" in e.lower() for e in exts)
    return {
        "file": os.path.basename(path),
        "tris": tris,
        "dims_m": dims,                     # x,y,z extent in file units (hopefully metres)
        "nodes": len(gltf.get("nodes", [])),
        "meshes": len(meshes),
        "materials": len(gltf.get("materials", [])),
        "images": len(images),
        "tex_sizes": tex_sizes,
        "max_tex": max((max(w, h) for w, h in tex_sizes), default=0),
        "draco": draco,
        "exts": exts,
    }


def _img_size(blob):
    if blob[:8] == b"\x89PNG\r\n\x1a\n":
        w, h = struct.unpack(">II", blob[16:24])
        return (w, h)
    if blob[:2] == b"\xff\xd8":              # JPEG: scan SOF markers
        i = 2
        while i < len(blob) - 9:
            if blob[i] != 0xFF:
                i += 1
                continue
            marker = blob[i + 1]
            if marker in (0xC0, 0xC1, 0xC2, 0xC3):
                h, w = struct.unpack(">HH", blob[i + 5:i + 9])
                return (w, h)
            seglen = struct.unpack(">H", blob[i + 2:i + 4])[0]
            i += 2 + seglen
    if blob[:4] == b"RIFF" and blob[8:12] == b"WEBP":     # WebP (normalize.py exports these)
        fmt = blob[12:16]
        if fmt == b"VP8 ":                   # lossy: 16-bit dims at offset 26 (14-bit each)
            w = struct.unpack("<H", blob[26:28])[0] & 0x3FFF
            h = struct.unpack("<H", blob[28:30])[0] & 0x3FFF
            return (w, h)
        if fmt == b"VP8L":                   # lossless: 14-bit dims packed after the 0x2F sig
            b0, b1, b2, b3 = blob[21], blob[22], blob[23], blob[24]
            w = ((b1 & 0x3F) << 8 | b0) + 1
            h = ((b3 & 0x0F) << 10 | b2 << 2 | (b1 & 0xC0) >> 6) + 1
            return (w, h)
        if fmt == b"VP8X":                   # extended: 24-bit dims at offset 24
            w = (blob[24] | blob[25] << 8 | blob[26] << 16) + 1
            h = (blob[27] | blob[28] << 8 | blob[29] << 16) + 1
            return (w, h)
    return None


def main():
    args = sys.argv[1:]
    as_json = "--json" in args
    args = [a for a in args if a != "--json"]
    files = []
    for a in args:
        if os.path.isdir(a):
            files += sorted(glob.glob(os.path.join(a, "**", "*.glb"), recursive=True))
        else:
            files.append(a)
    rows = []
    for f in files:
        try:
            rows.append(stat(f))
        except Exception as e:
            rows.append({"file": os.path.basename(f), "error": str(e)})
    if as_json:
        print(json.dumps(rows, indent=2))
        return
    hdr = f"{'file':38} {'tris':>7} {'dims (m) x,y,z':>22} {'mtl':>3} {'img':>3} {'maxtex':>6} draco"
    print(hdr)
    print("-" * len(hdr))
    for r in rows:
        if "error" in r:
            print(f"{r['file']:38} ERROR: {r['error']}")
            continue
        d = r["dims_m"]
        dstr = f"{d[0]:.2f},{d[1]:.2f},{d[2]:.2f}"
        print(f"{r['file']:38} {r['tris']:>7} {dstr:>22} {r['materials']:>3} "
              f"{r['images']:>3} {r['max_tex']:>6} {'YES' if r['draco'] else '-'}")


if __name__ == "__main__":
    main()