#!/usr/bin/env python3
"""Decode a `REGN` object (a stage's `MapPath`, `hidden/MiscBin.pak`).

A `REGN` object is a serialised C++ object graph with a trailing `POF0`
pointer-fixup table.  The fixup table is the primary evidence used here: it
names every word in the file that the retail loader turns into a pointer, so
the pointer graph does not have to be guessed.  The loader code that consumes
it is `sub_82465110` / `sub_82465138` / `sub_82465200` (see
`docs/re/structures/regn-map-grid.md`).

    ./regn_decode.py list   <MiscBin.pak>            # entries + magics
    ./regn_decode.py dump   <MiscBin.pak> [hash]     # one object's structure
    ./regn_decode.py verify <MiscBin.pak>            # the checks in the doc

Everything is big-endian.  All file offsets printed are absolute, i.e. the
stored pointer value plus 0x10 (the fixup base is chunk+0x10).
"""

import glob
import itertools
import math
import os
import random
import struct
import sys
import zlib

FIXUP_BASE = 0x10  # sub_82465138: the POF0 applier is called with r3 = chunk+16

# stride of each of the six sections, in bytes
STRIDE = (12, 96, 48, 8, 32, 4)
SECTION_NAME = (
    "vertices", "tetrahedra", "faces", "cell index", "cell items", "tet refs",
)
PAIRS = list(itertools.combinations(range(4), 2))


# ── IPFB archive ────────────────────────────────────────────────────────────

def pak_entries(pak_path):
    """[(name_hash, payload_bytes)] for every entry of an IPFB archive."""
    index = open(pak_path, "rb").read()
    if index[:4] != b"IPFB":
        raise SystemExit(f"{pak_path}: not an IPFB index")
    count, = struct.unpack_from(">I", index, 4)
    toc = [struct.unpack_from(">III", index, 0x10 + 12 * i) for i in range(count)]
    segs = sorted(glob.glob(os.path.splitext(pak_path)[0] + ".p[0-9][0-9]"))
    data = b"".join(open(s, "rb").read() for s in segs)
    out = []
    for name_hash, offset, comp_size in toc:
        stored = data[offset:offset + comp_size]
        if stored[:2] == b"Z1":
            stored = zlib.decompress(stored[10:])
        out.append((name_hash, stored))
    return out


# ── POF0 ────────────────────────────────────────────────────────────────────

def pof0_pointer_slots(blob):
    """Byte offsets of every word the retail loader relocates.

    Mirrors `sub_82465200` exactly: a run-length delta stream over word
    indices, 6 / 14 / 22 bits selected by the top two bits of the lead byte.
    """
    data_size, = struct.unpack_from(">I", blob, 4)
    table = data_size + 16
    if blob[table:table + 4] != b"POF0":
        raise SystemExit("no POF0 chunk at header[0x04]+16")
    size, = struct.unpack_from(">I", blob, table + 4)
    p, end, word = table + 16, table + 16 + size, 0
    slots = []
    while p < end:
        b = blob[p]
        if b == 0:
            break
        tag = b & 0xC0
        if tag == 0x40:
            word += b & 0x3F
            p += 1
        elif tag == 0x80:
            word += ((b & 0x3F) << 8) + blob[p + 1]
            p += 2
        elif tag == 0xC0:
            word += ((((b & 0x3F) << 8) + blob[p + 1]) << 8) + blob[p + 2]
            p += 3
        else:
            p += 1
            continue
        slots.append(FIXUP_BASE + 4 * word)
    return slots


# ── the object ──────────────────────────────────────────────────────────────

class Regn:
    def __init__(self, blob):
        if blob[:4] != b"REGN":
            raise SystemExit(f"not a REGN object (magic {blob[:4]!r})")
        self.b = blob
        self.data_size, = struct.unpack_from(">I", blob, 4)
        self.bbox_min = self.f3(0x10)
        self.bbox_max = self.f3(0x20)
        self.extent = self.f3(0x30)
        self.cell_size = self.f3(0x40)
        self.dims = [self.u32(0x50 + 4 * i) for i in range(3)]
        self.count = [self.u16(0x60 + 2 * i) for i in range(6)]
        # the six section pointers, resolved to absolute file offsets
        self.sec = [self.u32(0x70 + 4 * i) + FIXUP_BASE for i in range(6)]

    def u32(self, o):  return struct.unpack_from(">I", self.b, o)[0]
    def u16(self, o):  return struct.unpack_from(">H", self.b, o)[0]
    def f32(self, o):  return struct.unpack_from(">f", self.b, o)[0]
    def f3(self, o):   return tuple(self.f32(o + 4 * i) for i in range(3))
    def ptr(self, o):  return self.u32(o) + FIXUP_BASE

    # section 0 — vertex
    def vertex(self, i):
        return self.f3(self.sec[0] + 12 * i)

    # section 1 — tetrahedron
    def tet(self, i):
        b = self.sec[1] + 96 * i
        return {
            "sphere_centre": self.f3(b),
            "sphere_radius": self.f32(b + 12),
            "vertices":  [self.u16(b + 16 + 2 * k) for k in range(4)],
            "faces":     [self.u16(b + 24 + 2 * k) for k in range(4)],
            # six (face_i, face_j) pairs in itertools.combinations order
            "portal_cost": [self.f32(b + 32 + 8 * k) for k in range(6)],
            "portal_unk":  [self.f32(b + 36 + 8 * k) for k in range(6)],
            "index": self.u32(b + 80),
            "flags": self.b[b + 84:b + 88],
        }

    # section 2 — face (a plane plus its adjacency)
    def face(self, i):
        b = self.sec[2] + 48 * i
        return {
            "normal": self.f3(b),
            "d": self.f32(b + 12),
            "point": self.f3(b + 16),
            "one": self.f32(b + 28),
            "vertices": [self.u16(b + 32 + 2 * k) for k in range(3)],
            "index":  self.u16(b + 38),
            "tet":    [self.u16(b + 40), self.u16(b + 42)],   # 0xFFFF = hull
            "slot":   [self.u16(b + 44), self.u16(b + 46)],
        }

    # section 3/4/5 — the grid
    def cell_index(self, x, y, z):
        return (z * self.dims[1] + y) * self.dims[0] + x

    def cell(self, ci):
        """[tet index] for one cell — count/pointer, item record, ref array."""
        b = self.sec[3] + 8 * ci
        n, p = self.u32(b), self.u32(b + 4)
        if n == 0 or p == 0:
            return []
        item = p + FIXUP_BASE
        k = self.u32(item + 0x10)
        refs = self.ptr(item + 0x14)
        return [(self.ptr(refs + 4 * j) - self.sec[1]) // 96 for j in range(k)]

    def cell_bounds(self, ci):
        dx, dy = self.dims[0], self.dims[1]
        x, y, z = ci % dx, (ci // dx) % dy, ci // (dx * dy)
        lo = tuple(self.bbox_min[k] + self.cell_size[k] * (x, y, z)[k]
                  for k in range(3))
        return lo, tuple(lo[k] + self.cell_size[k] for k in range(3))


# ── commands ────────────────────────────────────────────────────────────────

def cmd_list(pak):
    for name_hash, blob in pak_entries(pak):
        magic = blob[:4].decode("latin1")
        magic = magic if magic.isprintable() else "?"
        print(f"{name_hash:08x}  {magic:4s}  {len(blob):9d}")


def cmd_dump(pak, want=None):
    for name_hash, blob in pak_entries(pak):
        if blob[:4] != b"REGN":
            continue
        if want is not None and name_hash != want:
            continue
        r = Regn(blob)
        print(f"── {name_hash:08x}  {len(blob)} bytes")
        print(f"   bbox     {r.bbox_min} .. {r.bbox_max}")
        print(f"   cell     {r.cell_size}   dims {r.dims}")
        print(f"   counts   {r.count}")
        for i in range(6):
            print(f"   sec{i} {SECTION_NAME[i]:<11s} @0x{r.sec[i]:06x}  "
                  f"{r.count[i]:6d} x {STRIDE[i]}")
        t = r.tet(0)
        print(f"   tet 0    verts {t['vertices']} faces {t['faces']} "
              f"r={t['sphere_radius']:.1f}")
        f0 = r.face(t["faces"][0])
        print(f"   face {t['faces'][0]:<4d} verts {f0['vertices']} "
              f"tets {f0['tet']} slots {f0['slot']}")
        occupied = sum(1 for ci in range(r.dims[0] * r.dims[1] * r.dims[2])
                       if r.cell(ci))
        print(f"   cells    {occupied} occupied of "
              f"{r.dims[0] * r.dims[1] * r.dims[2]}")
        if want is not None:
            return


def cmd_verify(pak):
    random.seed(0)
    objs = [(h, b) for h, b in pak_entries(pak) if b[:4] == b"REGN"]
    print(f"{len(objs)} REGN objects\n")
    hdr = ("object    ptr-slots        chain  face==3 verts   ctrl   "
           "sphere/cell  ctrl   portal cost")
    print(hdr)
    for name_hash, blob in objs:
        r = Regn(blob)
        slots = pof0_pointer_slots(blob)

        # 1. where the fixup table says the pointers are
        header = [o for o in slots if o < r.sec[0]]
        in3 = [o for o in slots if r.sec[3] <= o < r.sec[4]]
        in4 = [o for o in slots if r.sec[4] <= o < r.sec[5]]
        in5 = [o for o in slots if o >= r.sec[5]]
        others = len(slots) - len(header) - len(in3) - len(in4) - len(in5)
        shape = (header == [0x70 + 4 * i for i in range(6)]
                 and others == 0
                 and all((o - r.sec[3]) % 8 == 4 for o in in3)
                 and all((o - r.sec[4]) % 32 == 20 for o in in4)
                 and len(in5) == r.count[5]
                 and len(in4) == r.count[4] == len(in3))

        # 2. the cell -> item -> ref-array chain is contiguously packed
        total, chain = 0, True
        for i in range(r.count[4]):
            item = r.sec[4] + 32 * i
            if r.ptr(item + 0x14) != r.sec[5] + 4 * total:
                chain = False
            total += r.u32(item + 0x10)
        chain = chain and total == r.count[5]

        # 3. each of a tet's four faces passes through exactly 3 of its
        #    four vertices — against a random-face control
        good = ctrl = n = 0
        for i in range(r.count[1]):
            t = r.tet(i)
            V = [r.vertex(p) for p in t["vertices"]]
            scale = max(abs(x) for v in V for x in v) + 1
            for want_real in (True, False):
                for fi in (t["faces"] if want_real else
                           [random.randrange(r.count[2]) for _ in range(4)]):
                    fc = r.face(fi)
                    on = sum(1 for v in V
                             if abs(sum(fc["normal"][j] * v[j] for j in range(3))
                                    + fc["d"]) <= 1e-4 * scale)
                    if want_real:
                        good += on == 3
                        n += 1
                    else:
                        ctrl += on == 3

        # 4. a listed tet's bounding sphere reaches the listed cell —
        #    control is the same test with the cell axes transposed
        hit = thit = m = 0
        for ci in range(r.dims[0] * r.dims[1] * r.dims[2]):
            tets = r.cell(ci)
            if not tets:
                continue
            lo, hi = r.cell_bounds(ci)
            tlo, thi = tuple(reversed(lo)), tuple(reversed(hi))
            for ti in tets:
                t = r.tet(ti)
                c, rad = t["sphere_centre"], t["sphere_radius"]
                for (a, b), which in (((lo, hi), 0), ((tlo, thi), 1)):
                    q = 0.0
                    for j in range(3):
                        if c[j] < a[j]:
                            q += (a[j] - c[j]) ** 2
                        elif c[j] > b[j]:
                            q += (c[j] - b[j]) ** 2
                    if q <= rad * rad:
                        if which == 0:
                            hit += 1
                        else:
                            thit += 1
                m += 1

        # 5. the six portal costs are the face-centroid distances
        pc = pn = 0
        for i in range(r.count[1]):
            t = r.tet(i)
            V = [r.vertex(p) for p in t["vertices"]]
            omit = []
            for fi in t["faces"]:
                fv = set(r.face(fi)["vertices"])
                miss = [k for k in range(4) if t["vertices"][k] not in fv]
                omit.append(miss[0] if len(miss) == 1 else None)
            if any(o is None for o in omit):
                continue
            for k, (a, b) in enumerate(PAIRS):
                want = math.dist(V[omit[a]], V[omit[b]]) / 3
                pn += 1
                pc += abs(t["portal_cost"][k] - want) <= 1e-4 * max(1.0, want)

        print(f"{name_hash:08x}  {'ok' if shape else 'BAD':>9s}  "
              f"{'ok' if chain else 'BAD':>7s}  "
              f"{100 * good / n:6.2f}% {100 * ctrl / n:6.2f}%   "
              f"{100 * hit / m:6.2f}% {100 * thit / m:6.2f}%   "
              f"{100 * pc / pn:7.3f}%")


def main():
    if len(sys.argv) < 3:
        print(__doc__)
        raise SystemExit(2)
    cmd, pak = sys.argv[1], sys.argv[2]
    if cmd == "list":
        cmd_list(pak)
    elif cmd == "dump":
        want = int(sys.argv[3], 16) if len(sys.argv) > 3 else None
        cmd_dump(pak, want)
    elif cmd == "verify":
        cmd_verify(pak)
    else:
        print(__doc__)
        raise SystemExit(2)


if __name__ == "__main__":
    main()