Made of little things.

#!/usr/bin/env python3
"""Build game-ready reference data from the extracted original assets.

Dev-only tooling. Inputs come from extract_assets.py, hke.py, lightmap_uvs.py
and macromelt (see tools/README.md). Writes ref-out/game/:

  level.glb      visual level, Director Z-up units, with TEXCOORD_1 lightmap
                 UVs and mesh.extras.lightmap naming the lightmap texture
  car.glb, prize.glb, star.glb, arrow2.glb
  lightmaps/*.jpg, ui/*.{webp,jpg} (other bitmaps), sounds/*.{wav,mp3}

The bundle embeds all of this, so outputs are kept compact: GLBs share
duplicated textures and vertex data, UI bitmaps are WebP and one-shot sounds
are MP3 (needs ffmpeg). The engine sounds stay PCM WAV: they are queued back to
back and timed off their exact durations, which MP3 padding would change.
  collision.json Havok static bodies (world-space triangles + material)
  raymesh.bin    world-space triangles of every level model (the "3d2" copy the
                 wheels and camera raycast against): u32 counts, f32 xyz verts,
                 u32 indices, u16 model index per triangle; names in markers.json
  markers.json   W3D group nodes (checkpoints, prizes, camera targets) and
                 the level's cameras and model placements

Everything in ref-out/ is derived from copyrighted originals: local use only.
"""
import json
import os
import shutil
import struct
import subprocess

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
ASSETS = os.path.join(ROOT, "ref-out", "assets")
OUT = os.path.join(ROOT, "ref-out", "game")
MACROMELT = os.path.join(ROOT, "tools", "vendor", "macromelt", "target", "release", "macromelt")


# --- GLB helpers -------------------------------------------------------------

def read_glb(path):
    b = open(path, "rb").read()
    jlen = struct.unpack_from("<I", b, 12)[0]
    doc = json.loads(b[20:20 + jlen])
    off = 20 + jlen
    blen = struct.unpack_from("<I", b, off)[0]
    return doc, bytearray(b[off + 8:off + 8 + blen])


def write_glb(path, doc, binary):
    while len(binary) % 4:
        binary.append(0)
    doc["buffers"] = [{"byteLength": len(binary)}]
    j = json.dumps(doc, separators=(",", ":")).encode()
    j += b" " * (-len(j) % 4)
    total = 12 + 8 + len(j) + 8 + len(binary)
    with open(path, "wb") as f:
        f.write(struct.pack("<III", 0x46546C67, 2, total))
        f.write(struct.pack("<II", len(j), 0x4E4F534A) + j)
        f.write(struct.pack("<II", len(binary), 0x004E4942) + bytes(binary))


def dedup_glb(doc, binary):
    """Share byte-identical buffer views, images, textures and accessors.

    Materials, meshes and nodes are left alone (the game looks materials up by
    name). Returns the compacted binary.
    """
    # Drop accessors nothing references (e.g. replaced vertex data).
    used_acc = sorted({a for m in doc.get("meshes", []) for p in m["primitives"]
                       for a in [*p["attributes"].values(), *([p["indices"]] if "indices" in p else [])]})
    acc_keep = {a: i for i, a in enumerate(used_acc)}
    doc["accessors"] = [doc["accessors"][a] for a in used_acc]
    for m in doc.get("meshes", []):
        for p in m["primitives"]:
            p["attributes"] = {k: acc_keep[a] for k, a in p["attributes"].items()}
            if "indices" in p:
                p["indices"] = acc_keep[p["indices"]]

    def view_bytes(i):
        v = doc["bufferViews"][i]
        o = v.get("byteOffset", 0)
        return bytes(binary[o:o + v["byteLength"]])

    referenced = {a["bufferView"] for a in doc["accessors"]} | {im["bufferView"] for im in doc.get("images", [])}
    views, view_map = {}, {}
    for i, v in enumerate(doc["bufferViews"]):
        if i not in referenced:
            continue
        key = (view_bytes(i), v.get("byteStride"), v.get("target"))
        view_map[i] = views.setdefault(key, len(views))
    out = bytearray()
    new_views = []
    for (data, stride, target) in views:
        out += b"\0" * (-len(out) % 4)
        v = {"buffer": 0, "byteOffset": len(out), "byteLength": len(data)}
        if stride:
            v["byteStride"] = stride
        if target:
            v["target"] = target
        new_views.append(v)
        out += data
    doc["bufferViews"] = new_views

    def merge(items, remap_fields):
        seen, mapping, kept = {}, {}, []
        for i, item in enumerate(items):
            for f, m in remap_fields:
                if f in item:
                    item[f] = m[item[f]]
            key = json.dumps(item, sort_keys=True)
            if key not in seen:
                seen[key] = len(kept)
                kept.append(item)
            mapping[i] = seen[key]
        return kept, mapping

    doc["accessors"], acc_map = merge(doc.get("accessors", []), [("bufferView", view_map)])
    # Image names are only labels; identical bytes are the same picture.
    for im in doc.get("images", []):
        im.pop("name", None)
    doc["images"], img_map = merge(doc.get("images", []), [("bufferView", view_map)])
    doc["textures"], tex_map = merge(doc.get("textures", []), [("source", img_map)])
    for m in doc.get("materials", []):
        tex = m.get("pbrMetallicRoughness", {}).get("baseColorTexture")
        if tex:
            tex["index"] = tex_map[tex["index"]]
    for mesh in doc.get("meshes", []):
        for prim in mesh["primitives"]:
            prim["attributes"] = {k: acc_map[a] for k, a in prim["attributes"].items()}
            if "indices" in prim:
                prim["indices"] = acc_map[prim["indices"]]
    for key in ("images", "textures", "accessors"):
        if not doc[key]:
            del doc[key]
    return out


def compact_glb(src, dst):
    doc, binary = read_glb(src)
    write_glb(dst, doc, dedup_glb(doc, binary))


def accessor_floats(doc, binary, idx):
    acc = doc["accessors"][idx]
    view = doc["bufferViews"][acc["bufferView"]]
    width = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4}[acc["type"]]
    base = view.get("byteOffset", 0) + acc.get("byteOffset", 0)
    stride = view.get("byteStride", 4 * width)
    return [struct.unpack_from(f"<{width}f", binary, base + i * stride) for i in range(acc["count"])]


def add_accessor(doc, binary, values, kind="VEC"):
    """Append float vectors (kind "VEC") or triangle indices ("SCALAR")."""
    while len(binary) % 4:
        binary.append(0)
    off = len(binary)
    acc = {"bufferView": len(doc["bufferViews"]), "byteOffset": 0}
    if kind == "SCALAR":
        fmt, ctype = ("H", 5123) if max(values) < 65536 else ("I", 5125)
        binary += struct.pack(f"<{len(values)}{fmt}", *values)
        doc["bufferViews"].append({"buffer": 0, "byteOffset": off, "byteLength": len(binary) - off, "target": 34963})
        acc.update(componentType=ctype, count=len(values), type="SCALAR")
    else:
        width = len(values[0])
        for v in values:
            binary += struct.pack(f"<{width}f", *v)
        doc["bufferViews"].append({"buffer": 0, "byteOffset": off, "byteLength": 4 * width * len(values), "target": 34962})
        acc.update(componentType=5126, count=len(values), type=f"VEC{width}",
                   min=[min(v[c] for v in values) for c in range(width)],
                   max=[max(v[c] for v in values) for c in range(width)])
    del acc["byteOffset"]
    doc["accessors"].append(acc)
    return len(doc["accessors"]) - 1


# --- W3D materials (0x10): the GLB export keeps only the diffuse colour ---------

def w3d_materials(path):
    b = open(path, "rb").read()
    i = b.find(b"IFX") + 16
    out = {}
    while i + 8 <= len(b):
        kind, size = struct.unpack_from("<II", b, i)
        body = b[i + 8:i + 8 + size]
        if kind == 0xFFFFFF10:
            n = struct.unpack_from("<H", body, 0)[0]
            name = body[2:2 + n].decode("latin1")
            f = struct.unpack_from(f"<{(size - 2 - n - 4) // 4}f", body, 2 + n + 4)
            out[name] = dict(ambient=f[0:3], diffuse=f[4:7], specular=f[8:11], emissive=f[12:15],
                             shininess=f[16] if len(f) > 16 else 0, opacity=f[17] if len(f) > 17 else 1)
        i += 8 + ((size + 3) & ~3)
    return out


# --- W3D group nodes ----------------------------------------------------------

def w3d_nodes(path):
    """0x70 group, 0x72 model and 0x74 view nodes: name, parent, matrix."""
    b = open(path, "rb").read()
    i = b.find(b"IFX") + 16
    nodes = []
    while i + 8 <= len(b):
        kind, size = struct.unpack_from("<II", b, i)
        body = b[i + 8:i + 8 + size]
        if kind in (0xFFFFFF70, 0xFFFFFF72, 0xFFFFFF74):
            p = 0
            strs = []
            for _ in range(3):  # name, parent, user properties
                n = struct.unpack_from("<H", body, p)[0]
                strs.append(body[p + 2:p + 2 + n].decode("latin1"))
                p += 2 + n
            matrix = list(struct.unpack_from("<16f", body, p))
            nodes.append(dict(kind={0x70: "group", 0x72: "model", 0x74: "view"}[kind & 0xFF],
                              name=strs[0], parent=strs[1], user=strs[2], matrix=matrix))
        i += 8 + ((size + 3) & ~3)
    return nodes


def main():
    os.makedirs(os.path.join(OUT, "lightmaps"), exist_ok=True)
    geom = {e["name"]: e for e in json.load(open(os.path.join(ASSETS, "geom.json")))}
    lm = json.load(open(os.path.join(ASSETS, "lightmap_uvs.json")))
    doc, binary = read_glb(os.path.join(ASSETS, "glb", "3d.glb"))

    used_maps = set()
    for mesh in doc["meshes"]:
        name = mesh["name"].removeprefix("mesh:")
        info = lm.get(name)
        if not info or not info.get("lightmap"):
            continue
        g = geom[name]
        # lightMapManager.applyUVs sets meshDeform.mesh[n].textureLayer[2] per vertex index;
        # the blocks are contiguous per submesh, in order, like geom.json's vertices.
        flat = [uv for block in info["uvs"] for uv in block]
        assert len(flat) == len(g["positions"]), name
        assert len(mesh["primitives"]) == len(set(g["submesh_of_face"])), name
        # The GLB export reorders vertices, and seam vertices share position and normal
        # but not lightmap UV, so rebuild each primitive from the decoded mesh by index.
        texcoords = g["texcoords"]
        for k, prim in enumerate(mesh["primitives"]):
            faces = [f for f, s in zip(g["faces"], g["submesh_of_face"]) if s == k]
            used = sorted({i for f in faces for i in f})
            local = {v: j for j, v in enumerate(used)}
            attrs = prim["attributes"]
            attrs["POSITION"] = add_accessor(doc, binary, [g["positions"][i] for i in used])
            attrs["NORMAL"] = add_accessor(doc, binary, [g["normals"][i] for i in used])
            if "TEXCOORD_0" in attrs:
                # glTF's v runs down the image; Director's runs up.
                attrs["TEXCOORD_0"] = add_accessor(doc, binary, [(texcoords[i][0], 1 - texcoords[i][1]) for i in used])
            attrs["TEXCOORD_1"] = add_accessor(doc, binary, [tuple(flat[i]) for i in used])
            prim["indices"] = add_accessor(doc, binary, [local[i] for f in faces for i in f], "SCALAR")
        mesh.setdefault("extras", {})["lightmap"] = info["lightmap"] + "LightingMap"
        used_maps.add(info["lightmap"] + "LightingMap")
    write_glb(os.path.join(OUT, "level.glb"), doc, dedup_glb(doc, binary))
    print(f"level.glb: lightmaps {sorted(used_maps)}")

    # Lightmaps are used as rgba8880 (no alpha): JPEG keeps the bundle small.
    from PIL import Image
    for name in used_maps:
        Image.open(os.path.join(ASSETS, "bitmaps", f"{name}.png")).convert("RGB").save(
            os.path.join(OUT, "lightmaps", f"{name}.jpg"), quality=92)

    # UI bitmaps and sounds (lightmaps are copied above).
    for sub in ("ui", "sounds"):
        os.makedirs(os.path.join(OUT, sub), exist_ok=True)
    # Full-screen backdrops: their alpha plane is broken, the RGB is complete.
    opaque = {"Title-v8", "Help-v5", "Background"}
    for f in os.listdir(os.path.join(ASSETS, "bitmaps")):
        stem = f.rsplit(".", 1)[0]
        if "LightingMap" in f:
            continue
        if stem in opaque:
            Image.open(os.path.join(ASSETS, "bitmaps", f)).convert("RGB").save(os.path.join(OUT, "ui", f"{stem}.jpg"), quality=90)
        else:
            # Lossy colour, exact alpha: the button art keeps its hard edges.
            Image.open(os.path.join(ASSETS, "bitmaps", f)).save(
                os.path.join(OUT, "ui", f"{stem}.webp"), quality=90, alpha_quality=100, method=6)
    pcm_sounds = {"CarStart", "CarLoop", "CarEnd"}
    for f in os.listdir(os.path.join(ASSETS, "sounds")):
        stem, ext = f.rsplit(".", 1)
        src = os.path.join(ASSETS, "sounds", f)
        if ext == "wav" and stem not in pcm_sounds:
            subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", src, "-c:a", "libmp3lame", "-b:a", "32k",
                            os.path.join(OUT, "sounds", f"{stem}.mp3")], check=True)
        else:
            shutil.copy(src, os.path.join(OUT, "sounds", f))

    for src, dst in (("holypolycrap", "car"), ("prize", "prize"), ("star", "star"), ("arrow2", "arrow2")):
        compact_glb(os.path.join(ASSETS, "glb", f"{src}.glb"), os.path.join(OUT, f"{dst}.glb"))

    havok = json.load(open(os.path.join(ASSETS, "havok.json")))
    collision = dict(gravity=havok["gravity_world"], scale=havok["scale"], bodies=[
        dict(name=b["name"], convex=b["convex"], friction=b["friction"], restitution=b["restitution"],
             verts=[round(c, 3) for v in b["world_verts"] for c in v], faces=[i for f in b["faces"] for i in f])
        for b in havok["bodies"]])
    json.dump(collision, open(os.path.join(OUT, "collision.json"), "w"), separators=(",", ":"))

    # Ray mesh: every level model in world space, as in the duplicated "3d2" member.
    names, verts, tris, owner = [], [], [], []
    for gi, g in enumerate(geom.values()):
        m = g["node_matrix"]  # row-vector convention, translation in 12..14
        base = len(verts) // 3
        for x, y, z in g["positions"]:
            verts += [x * m[0] + y * m[4] + z * m[8] + m[12], x * m[1] + y * m[5] + z * m[9] + m[13],
                      x * m[2] + y * m[6] + z * m[10] + m[14]]
        for f in g["faces"]:
            tris += [base + f[0], base + f[1], base + f[2]]
            owner.append(gi)
        names.append(g["name"])
    with open(os.path.join(OUT, "raymesh.bin"), "wb") as f:
        f.write(struct.pack("<II", len(verts) // 3, len(tris) // 3))
        f.write(struct.pack(f"<{len(verts)}f", *verts))
        f.write(struct.pack(f"<{len(tris)}I", *tris))
        f.write(struct.pack(f"<{len(owner)}H", *owner))
    car = geom["car"]
    hull = dict(name="car", node_matrix=car["node_matrix"], positions=car["positions"])
    print(f"raymesh.bin: {len(verts) // 3} verts, {len(tris) // 3} tris")

    materials = {f: w3d_materials(os.path.join(ASSETS, "3d", f"{f}.w3d")) for f in ("3d", "holypolycrap", "star", "prize", "arrow2")}
    json.dump(materials, open(os.path.join(OUT, "materials.json"), "w"), indent=1)

    markers = dict(raymesh_models=names, car_hull=hull, level=w3d_nodes(os.path.join(ASSETS, "3d", "3d.w3d")),
                   car=w3d_nodes(os.path.join(ASSETS, "3d", "holypolycrap.w3d")))
    json.dump(markers, open(os.path.join(OUT, "markers.json"), "w"), indent=1)
    print(f"collision.json: {len(collision['bodies'])} bodies; markers.json: {len(markers['level'])} level nodes")


if __name__ == "__main__":
    if not os.path.exists(os.path.join(ASSETS, "geom.json")):
        subprocess.run([MACROMELT, os.path.join(ASSETS, "3d", "3d.w3d"), "-q", "--dump-geom",
                        os.path.join(ASSETS, "geom.json")], check=True)
    main()