#!/usr/bin/env python3
"""Solve the BASE of every register-relative name block, and print the schema it names.

A loader that reads a table by field name usually does not take each string's
address directly -- it keeps one base pointer and emits `addi rX, rBASE, -N` per
name.  No static xref sees that, so those blocks look unreferenced.

This solves the base from the DISPLACEMENT SET alone: for a group of `addi`s
sharing a source register, try every candidate base implied by (string address -
displacement) and keep the one that lands the most displacements on a string
start.  Control: it must recover `r30 = 0x82088F94` for the unit loader
`sub_82341A20` with no prior knowledge.

Regenerates docs/re/data/name-block-bases.txt.
"""
import sys, os, glob, collections, bisect

def _resolve_db():
    """`$SYLPHEED_DB`, else `$SYLPH_XEXDB`, else `<repo root>/sylpheed.db`.

    `SYLPHEED_DB` is the documented contract -- `docker/decoder/sylph-decoder` sets
    it and `docs/agents/CONTAINER-NOTES.md` lists it -- but no reader honoured it:
    this script hardcoded `/work/xenia-rs/sylpheed.db`, a path that has not
    existed anywhere since `/work` became a clone, so it failed on every machine.
    """
    import os, pathlib
    for var in ('SYLPHEED_DB', 'SYLPH_XEXDB'):
        v = os.environ.get(var)
        if v:
            p = pathlib.Path(v).expanduser()
            if not p.is_file():
                sys.exit(f'${var} is set but is not a file: {p}')
            return str(p)
    p = pathlib.Path(__file__).resolve().parents[2] / 'sylpheed.db'
    if p.is_file():
        return str(p)
    sys.exit(f'no database: set $SYLPHEED_DB, or build {p} with '
             'sylph-xexdb dis <xex|iso> --db sylpheed.db --analyze sql --quiet')


MIN_GROUP = 8          # displacements needed before a group is worth solving
MIN_RESOLVED = 12      # report a function only if the base explains this many
LOW, HIGH = 0x82000000, 0x82400000

def main():
    import duckdb
    con = duckdb.connect(_resolve_db(), read_only=True)
    S = set(a for (a,) in con.execute("SELECT address FROM strings").fetchall())
    txt = dict(con.execute("SELECT address, content FROM strings").fetchall())
    funcs = con.execute(
        "SELECT address,end_address,name FROM functions ORDER BY address").fetchall()
    starts = [f[0] for f in funcs]
    ins = con.execute(
        "SELECT address,operands FROM instructions WHERE mnemonic='addi' ORDER BY address"
    ).fetchall()

    print("# Register-relative name blocks: the base solved from the displacements")
    print("# Regenerate: python3 tools/re-capture/name_block_bases.py")
    print("# See docs/re/structures/player-tuning-tables.md")
    print("\n   strings in the image: %d   addi instructions: %d" % (len(S), len(ins)))

    byf = collections.defaultdict(lambda: collections.defaultdict(list))
    order = collections.defaultdict(lambda: collections.defaultdict(list))
    for a, o in ins:
        i = bisect.bisect_right(starts, a) - 1
        if i < 0 or funcs[i][1] <= a:
            continue
        p = [x.strip() for x in o.split(',')]
        if len(p) != 3 or p[0] == p[1]:
            continue
        try:
            d = int(p[2], 0)
        except ValueError:
            continue
        byf[funcs[i][2]][p[1]].append(d)
        order[funcs[i][2]][p[1]].append((a, d))

    rows = []
    for name in sorted(byf):
        for reg in sorted(byf[name]):
            D = sorted(set(byf[name][reg]))
            # r1 is the stack pointer.  r0 is NOT a base register at all: in
            # `addi rD, r0, N` the RA slot reads as literal zero, so the form is
            # `li rD, N` and the displacements are plain immediates.  Left in, the
            # VMX save/restore helper pair sub_825F2CF0 / sub_825F2F88 (72 x
            # `addi r11, r0, -N`, the vector spill offsets) solved a base and
            # "named" 30 strings, 97% of them disc names -- a pure artefact.
            if len(D) < MIN_GROUP or reg in ('r0', 'r1'):
                continue
            # Vote: every (string, displacement) pair implies one candidate base.
            # The true base collects a vote from each name it explains, so it wins
            # outright -- taking candidates from ONE displacement misses it.
            cand = collections.Counter()
            for A in S:
                for d in D:
                    B = A - d
                    if LOW <= B <= HIGH:
                        cand[B] += 1
            best = None
            # sorted, not most_common(): a Counter's tie order varies per run
            top = sorted(cand.items(), key=lambda kv: (-kv[1], kv[0]))[:40]
            for B, _v in top:
                tot = sum(1 for d in D if (B + d) in S)
                if best is None or tot > best[1] or (tot == best[1] and B < best[0]):
                    best = (B, tot)
            if best and best[1] >= MIN_RESOLVED:
                rows.append((name, reg, best[0], best[1], len(D)))

    rows.sort(key=lambda r: (-r[3], r[0]))
    print("\n## %d functions read a name block through a base register" % len(rows))
    print("   %-18s %-5s %-12s %s" % ('function', 'reg', 'base', 'resolved / displacements'))
    for name, reg, B, tot, n in rows:
        print("   %-18s %-5s 0x%08X   %d / %d" % (name, reg, B, tot, n))

    # Does the function itself build the solved base into the solved register?
    # Simulate `lis/addis rD, r0, HI`, `addi rD, rA, N` and `or rD, rA, rA`
    # forward through the body; a hit is POSITIVE confirmation of the base.
    # (A miss is not a refutation: a callee-saved base is often materialised in
    # the caller or loaded from memory, which is why the solver exists at all.)
    want = collections.defaultdict(set)
    for name, reg, B, tot, n in rows:
        want[name].add((reg, B))
    sim = con.execute(
        "SELECT address,mnemonic,operands FROM instructions "
        "WHERE mnemonic IN ('addi','addis','lis','or') ORDER BY address").fetchall()
    body = collections.defaultdict(list)
    for a, m, o in sim:
        i = bisect.bisect_right(starts, a) - 1
        if i < 0 or funcs[i][1] <= a:
            continue
        if funcs[i][2] in want:
            body[funcs[i][2]].append((m, o))
    confirmed = set()
    for name, pairs in sorted(body.items()):
        targets = want[name]
        val = {}
        for m, o in pairs:
            q = [x.strip() for x in o.split(',')]
            if len(q) != 3:
                continue
            if m in ('addis', 'lis') and q[1] == 'r0':
                try:
                    val[q[0]] = (int(q[2], 0) << 16) & 0xFFFFFFFF
                except ValueError:
                    val.pop(q[0], None)
            elif m == 'addi':
                try:
                    d = int(q[2], 0)
                except ValueError:
                    val.pop(q[0], None)
                    continue
                if q[1] == 'r0':
                    val[q[0]] = d & 0xFFFFFFFF
                elif q[1] in val:
                    val[q[0]] = (val[q[1]] + d) & 0xFFFFFFFF
                else:
                    val.pop(q[0], None)
            elif m == 'or' and q[1] == q[2]:
                if q[1] in val:
                    val[q[0]] = val[q[1]]
                else:
                    val.pop(q[0], None)
            for reg, B in sorted(targets):
                if val.get(reg) == B:
                    confirmed.add((name, reg, B))

    round_rows = [r for r in rows if (r[2] & 0xFFFF) == 0]
    solved_rows = [r for r in rows if (r[2] & 0xFFFF) != 0]
    ok = lambda rs: sum(1 for r in rs if (r[0], r[1], r[2]) in confirmed)
    print("\n## Is the base built by the function itself?  (positive test only)")
    print("   rows whose function materialises the solved base into the solved"
          " register: %d / %d" % (ok(rows), len(rows)))
    print("      bases on a 64K boundary                        : %d / %d"
          % (ok(round_rows), len(round_rows)))
    print("      bases with a non-zero low half                 : %d / %d"
          % (ok(solved_rows), len(solved_rows)))
    print("   A 64K-boundary base is `addis rX, r0, 0xHHHH` written out in full, so it"
          " confirms\n   directly; it is NOT the low-confidence class it was first"
          " called.  A non-zero low\n   half usually means the base was built in the"
          " caller or loaded from memory, which\n   the simulation cannot see -- a miss"
          " there is silence, not a refutation.")
    rb = collections.Counter(r[2] for r in round_rows)
    print("\n   64K-boundary bases by value:")
    for B, c in sorted(rb.items(), key=lambda kv: (-kv[1], kv[0])):
        print("      0x%08X  x%d rows" % (B, c))
    print("   \u26a0 60 of the 0x820B0000 rows are one function emitted 60 times:"
          " exactly 491\n     instructions each, an identical sequence of 41 string"
          " addresses, differing only\n     in two global data pointers.  That is"
          " DUPLICATION, not error.")

    # Is a row a DATA-TABLE schema or engine/XDK vocabulary?  The objective test
    # is whether its names are IDXD record/field names on the disc.
    sys.path.insert(0, HERE) if False else None
    from unit_substructures import pak_entries
    import unitgroup as _U
    disc = set()
    for pk in sorted(glob.glob('/work/sylph_extract/**/*.pak', recursive=True)):
        for _h, b in pak_entries(pk):
            if b[:4] != b'IDXD':
                continue
            try:
                rs = _U.parse(b)
            except Exception:
                continue
            for r in rs:
                disc.add(r['squadron'])
                for _t, n, _v in r['fields']:
                    if n:
                        disc.add(n)
    print("\n## Is the row a DATA-TABLE schema?  (names that are IDXD record/field"
          " names on the disc)")
    print("   distinct IDXD record+field names disc-wide: %d" % len(disc))
    scored = []
    for name, reg, B, tot, n in rows:
        nm, seen = [], set()
        for a, d in order[name][reg]:
            t = B + d
            if t in txt and t not in seen:
                seen.add(t)
                nm.append(txt[t])
        if not nm:
            continue
        hit = sum(1 for x in nm if x in disc)
        scored.append((hit / len(nm), hit, len(nm), name, reg, B))
    tab = [r for r in scored if r[0] >= 0.5 and r[2] >= 8]
    print("   rows that are >=50%% disc names and >=8 names: %d / %d"
          % (len(tab), len(scored)))
    for r in sorted(tab, key=lambda r: (-r[2], r[3])):
        print("   %-14s %-4s 0x%08X %4d names %3.0f%% disc"
              % (r[3], r[4], r[5], r[2], 100 * r[0]))

    print("\n## The schema each one names, in code order")
    for name, reg, B, tot, n in rows:
        seen, names = set(), []
        for a, d in order[name][reg]:
            t = B + d
            if t in txt and t not in seen:
                seen.add(t)
                names.append(txt[t])
        if not names:
            continue
        print("\n   %s  (%s, base 0x%08X) names %d:" % (name, reg, B, len(names)))
        line = "     "
        for nm in names:
            if len(line) + len(nm) > 96:
                print(line)
                line = "     "
            line += nm + "  "
        if line.strip():
            print(line)

if __name__ == "__main__":
    main()