#!/usr/bin/env python3
"""MSVC RTTI walker for the SOTS1 exe: type descriptors -> COLs -> vftables.

Why this exists
---------------
`findings/objects/00-inventory.md` has 1,924 type-descriptor *names* but no
addresses, so a recovered function can be named "a serializer" but not "*whose*
serializer".  Every serializable class in this binary reaches its `Read`/`Write`
through an `IStreamable` sub-vftable, and MSVC puts a Complete Object Locator at
`vftable[-1]`.  Walking

    type descriptor  <-  COL.pTypeDescriptor  <-  vftable[-1]

gives, for every vftable in the image: the owning class name, the sub-object
offset (`COL.offset` -- the this-adjustment that turns a decompiled offset into
an absolute member offset), and the slot functions.

Structures (32-bit MSVC):
    TypeDescriptor  { void* pVFTable; void* spare; char name[]; }   name at +8
    COL             { u32 signature; u32 offset; u32 cdOffset;
                      TypeDescriptor* pTypeDescriptor;
                      ClassHierarchyDescriptor* pClassDescriptor; }  20 bytes
    vftable[-1] = COL*

Usage:
    uv run python3 tools/rtti_map.py build      # -> dumps/rtti.json
    uv run python3 tools/rtti_map.py show ObservedTech
    uv run python3 tools/rtti_map.py vftable 0x00a2439c
    uv run python3 tools/rtti_map.py stats
"""
import json
import os
import re
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.dirname(HERE)
EXE = os.path.join(REPO, "dumps", "sots.exe")
FUNCS = os.path.join(REPO, "dumps", "functions.json")
RTTI = os.path.join(REPO, "dumps", "rtti.json")


# ----------------------------------------------------------------- PE image
class Image:
    """Whole-image VA reader (all sections, not just executable ones)."""

    def __init__(self, path=EXE):
        data = open(path, "rb").read()
        pe = struct.unpack_from("<I", data, 0x3C)[0]
        assert data[pe:pe + 4] == b"PE\0\0", "not a PE"
        nsec = struct.unpack_from("<H", data, pe + 6)[0]
        optsz = struct.unpack_from("<H", data, pe + 20)[0]
        self.base = struct.unpack_from("<I", data, pe + 24 + 28)[0]
        self.secs = []          # (va, size, bytes, name, exec)
        off = pe + 24 + optsz
        for k in range(nsec):
            s = off + k * 40
            name = data[s:s + 8].rstrip(b"\0").decode("latin1")
            vsize, vaddr, rsize, raddr = struct.unpack_from("<IIII", data, s + 8)
            chars = struct.unpack_from("<I", data, s + 36)[0]
            n = min(vsize, rsize) if vsize else rsize
            self.secs.append((self.base + vaddr, n, data[raddr:raddr + n],
                              name, bool(chars & 0x20000000)))
        self.lo = min(s[0] for s in self.secs)
        self.hi = max(s[0] + s[1] for s in self.secs)
        self.text = next(s for s in self.secs if s[4])

    def find(self, va):
        for sva, n, buf, name, ex in self.secs:
            if sva <= va < sva + n:
                return buf, va - sva, name
        return None, 0, None

    def u32(self, va):
        buf, o, _ = self.find(va)
        if buf is None or o + 4 > len(buf):
            return None
        return struct.unpack_from("<I", buf, o)[0]

    def i32(self, va):
        buf, o, _ = self.find(va)
        if buf is None or o + 4 > len(buf):
            return None
        return struct.unpack_from("<i", buf, o)[0]

    def cstr(self, va, maxn=512):
        buf, o, _ = self.find(va)
        if buf is None:
            return None
        e = buf.find(b"\0", o, o + maxn)
        if e < 0:
            return None
        try:
            return buf[o:e].decode("latin1")
        except UnicodeDecodeError:
            return None

    def in_image(self, va):
        return va is not None and self.lo <= va < self.hi

    def in_text(self, va):
        return va is not None and self.text[0] <= va < self.text[0] + self.text[1]


# ------------------------------------------------------------------ demangle
def demangle(mangled: str) -> str:
    """`.?AVObservedTech@Game@@` -> `Game::ObservedTech`.

    Templates keep their raw inner text (`?$StreamableHelper@...`) because a
    full MSVC template demangler is not needed here -- only a stable key.
    """
    m = mangled
    for p in (".?AV", ".?AU", ".?AW", ".?AT"):
        if m.startswith(p):
            m = m[len(p):]
            break
    if m.endswith("@@"):
        m = m[:-2]
    # A template argument list can itself contain '@'; split only the trailing
    # namespace chain, which is everything after the outermost template body.
    if m.startswith("?$"):
        return m                        # template: keep raw
    parts = m.split("@")
    return "::".join(reversed([p for p in parts if p]))


# --------------------------------------------------------------------- build
def build():
    img = Image()
    funcs = json.load(open(FUNCS))
    fstarts = {int(k, 16) for k in funcs}
    fnames = {int(k, 16): v[0] for k, v in funcs.items()}

    # 1. type descriptors: the mangled name lives at TD+8
    tds = {}                            # td_va -> mangled
    pat = re.compile(rb"\.\?A[VUWT][\x20-\x7e]{0,300}?@@\x00")
    for sva, n, buf, name, ex in img.secs:
        if ex:
            continue
        for m in pat.finditer(buf):
            nva = sva + m.start()
            td = nva - 8
            if img.in_image(td):
                tds[td] = m.group()[:-1].decode("latin1")

    # 2. COLs: any aligned dword equal to a TD address, at COL+0xC
    cols = {}                           # col_va -> dict
    for sva, n, buf, name, ex in img.secs:
        if ex:
            continue
        for o in range(0, n - 3, 4):
            v = struct.unpack_from("<I", buf, o)[0]
            if v not in tds:
                continue
            col = sva + o - 0xC
            sig = img.u32(col)
            if sig not in (0, 1):
                continue
            off = img.u32(col + 4)
            cd = img.u32(col + 8)
            chd = img.u32(col + 0x10)
            if off is None or off > 0x10000 or cd is None or cd > 0x10000:
                continue
            if not img.in_image(chd):
                continue
            # Class Hierarchy Descriptor -> base class list.  This is the only
            # exact test for "is this class an IStreamable?": a 3-slot vftable
            # on its own also matches TacAISquadRule_*, the row parsers and any
            # other class that happens to have three virtuals.
            bases = []
            nb = img.u32(chd + 8)
            arr = img.u32(chd + 0xC)
            if nb and nb < 64 and img.in_image(arr):
                for b in range(nb):
                    bcd = img.u32(arr + 4 * b)
                    if not img.in_image(bcd):
                        break
                    btd = img.u32(bcd)
                    if btd in tds:
                        bases.append(demangle(tds[btd]))
            cols[col] = {"td": v, "offset": off, "cdOffset": cd,
                         "mangled": tds[v], "name": demangle(tds[v]),
                         "bases": bases}

    # 3. vftables: any aligned dword equal to a COL address; table starts at +4
    vfts = {}
    for sva, n, buf, name, ex in img.secs:
        if ex:
            continue
        for o in range(0, n - 3, 4):
            v = struct.unpack_from("<I", buf, o)[0]
            if v not in cols:
                continue
            vf = sva + o + 4
            slots = []
            a = vf
            while True:
                p = img.u32(a)
                if p is None or p not in fstarts:
                    break
                slots.append(p)
                a += 4
                if len(slots) > 400:
                    break
            if not slots:
                continue
            c = cols[v]
            vfts[vf] = {"col": v, "class": c["name"], "mangled": c["mangled"],
                        "offset": c["offset"], "slots": slots,
                        "bases": c["bases"],
                        "slotNames": [fnames.get(s, "") for s in slots]}

    out = {"typeDescriptors": {hex(k): v for k, v in tds.items()},
           "cols": {hex(k): v for k, v in cols.items()},
           "vftables": {hex(k): v for k, v in vfts.items()}}
    with open(RTTI, "w") as fh:
        json.dump(out, fh)
    print(f"type descriptors : {len(tds)}")
    print(f"COLs             : {len(cols)}")
    print(f"vftables         : {len(vfts)}")
    ns = {}
    for v in vfts.values():
        ns[v["class"].split("::")[0]] = ns.get(v["class"].split("::")[0], 0) + 1
    print("vftables by top-level namespace:",
          ", ".join(f"{k}={v}" for k, v in
                    sorted(ns.items(), key=lambda x: -x[1])[:6]))


def load():
    with open(RTTI) as fh:
        r = json.load(fh)
    return ({int(k, 16): v for k, v in r["typeDescriptors"].items()},
            {int(k, 16): v for k, v in r["cols"].items()},
            {int(k, 16): v for k, v in r["vftables"].items()})


def main():
    if len(sys.argv) < 2 or sys.argv[1] == "build":
        return build()
    cmd = sys.argv[1]
    tds, cols, vfts = load()
    if cmd == "show":
        pat = sys.argv[2]
        for va, v in sorted(vfts.items()):
            if pat.lower() in v["class"].lower():
                print(f"vftable 0x{va:08x}  COL 0x{v['col']:08x}  "
                      f"offset +0x{v['offset']:x}  {v['class']}")
                for i, (s, nm) in enumerate(zip(v["slots"], v["slotNames"])):
                    print(f"    [{i}] 0x{s:08x}  {nm}")
    elif cmd == "vftable":
        va = int(sys.argv[2], 0)
        v = vfts.get(va)
        print(json.dumps(v, indent=2) if v else "not a vftable")
    elif cmd == "stats":
        print(f"tds={len(tds)} cols={len(cols)} vftables={len(vfts)}")
    return 0


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