#!/usr/bin/env python3
"""Discover the runtime class that carries the parsed `UN_*` unit definitions.

No vtable VA is assumed. Procedure:
  1. locate every disc unit-ID string in the snapshot;
  2. find every 4-byte-aligned word in RAM that points at (string_va - d) for a
     few plausible name-record deltas;
  3. for each such pointer site P and each small back-offset k, read the word at
     P-k and tally it across *distinct* unit IDs.
A word that appears at the same (d, k) for many different units is that class's
vtable pointer, and k is the ID-pointer offset inside the object.
"""
import os
import re
import struct
import sys
from collections import defaultdict

sys.path.insert(0, "/home/fabi/RE - Project Sylpheed/sylpheed-reborn/tools/re-capture")
import gmem  # noqa: E402

DELTAS = [0x00, 0x10, 0x08, 0x04, 0x0C, 0x14, 0x18, 0x20]
BACK = list(range(0, 0x60, 4))


def load_ids(path):
    ids = []
    for line in open(path):
        if line.startswith("REC "):
            p = line.split()
            if p[3] == "Generic" and p[4].startswith("UN_"):
                ids.append(p[4])
    return sorted(set(ids))


def scan_strings(f, size, ids):
    """{id: [va,...]} — every NUL-terminated occurrence of each id string."""
    pats = {i: (i.encode() + b"\0") for i in ids}
    rx = re.compile(b"|".join(re.escape(p) for p in pats.values()))
    out = defaultdict(list)
    for a, b in gmem.extents(f.fileno(), size):
        f.seek(a)
        blob = f.read(b - a)
        for m in rx.finditer(blob):
            off = a + m.start()
            va = gmem.primary_va(off)
            if va is not None:
                out[m.group(0)[:-1].decode()].append(va)
    return out


def scan_pointers(f, size, targets):
    """{target_va: [site_va,...]} for every aligned word equal to a target."""
    want = {struct.pack(">I", t): t for t in targets}
    rx = re.compile(b"|".join(re.escape(k) for k in want))
    out = defaultdict(list)
    for a, b in gmem.extents(f.fileno(), size):
        f.seek(a)
        blob = f.read(b - a)
        for m in rx.finditer(blob):
            off = a + m.start()
            if off % 4:
                continue
            va = gmem.primary_va(off)
            if va is not None:
                out[want[m.group(0)]].append(va)
    return out


def main():
    ids = load_ids(sys.argv[1])
    path = gmem.mem_path()
    size = os.path.getsize(path)
    with open(path, "rb") as f:
        strs = scan_strings(f, size, ids)
        print(f"# {len(ids)} disc unit IDs, {len(strs)} found in RAM "
              f"({sum(len(v) for v in strs.values())} occurrences)", file=sys.stderr)

        # every (id, string_va - delta) we want pointers to
        targets, owner = set(), defaultdict(set)
        for uid, vas in strs.items():
            for va in vas:
                for d in DELTAS:
                    t = va - d
                    targets.add(t)
                    owner[t].add((uid, d))
        ptrs = scan_pointers(f, size, targets)
        print(f"# {len(ptrs)} of {len(targets)} candidate targets are pointed at",
              file=sys.stderr)

        # tally (delta, back-offset, word) over distinct unit IDs
        tally = defaultdict(set)
        sites = defaultdict(list)
        for tgt, plist in ptrs.items():
            for uid, d in owner[tgt]:
                for p in plist:
                    for k in BACK:
                        base = p - k
                        try:
                            f.seek(gmem.va_to_off(base))
                        except ValueError:
                            continue
                        w = f.read(4)
                        if len(w) < 4:
                            continue
                        (vt,) = struct.unpack(">I", w)
                        if 0x82000000 <= vt < 0x83000000:
                            tally[(d, k, vt)].add(uid)
                            sites[(d, k, vt)].append((base, uid))

    rank = sorted(tally.items(), key=lambda kv: -len(kv[1]))
    print(f"{'delta':>6} {'idoff':>6} {'word':>10} {'#ids':>5}")
    for (d, k, vt), s in rank[:25]:
        print(f"{d:#6x} {k:#6x} {vt:#010x} {len(s):5d}   e.g. {sorted(s)[:2]}")
    if rank:
        best = rank[0][0]
        print("\n# object bases for the top candidate (first 20):", file=sys.stderr)
        for base, uid in sorted(set(sites[best]))[:20]:
            print(f"#   {base:#010x}  {uid}", file=sys.stderr)


if __name__ == "__main__":
    main()