#!/usr/bin/env python3
"""Solve the runtime `Weapon`/`Shell` struct layouts against disc ground truth,
then read the fields the disc leaves defaulted.

The game parses every `weapon\\Weapon_*.tbl` IDXD record into two C++ objects, a
`Weapon` and its `Shell`. Each class lives in its own contiguous array in guest
RAM and is identifiable by its vtable pointer. Fields the IDXD omits (the
Route-B "defaulted" list) still hold their real values in those objects — put
there by the constructor before the parser overwrites what the disc supplies.

Method (no guessing):
  1. read every disc record's explicitly-valued fields, split per sub-record
     (`sylpheed-formats --example idxd_tokens`);
  2. read every runtime object out of the live emulator (`gmem`), keyed by the
     object's own ID string;
  3. for each (field, byte-offset, encoding) triple, count how many weapons the
     offset reproduces and how many it contradicts. An offset that agrees on
     many and contradicts none is a confirmed binding;
  4. print the map, then the values at those offsets for the weapons whose disc
     record omits the field.

Step 3 is the whole safeguard: a wrong offset cannot silently agree with ~100
independent records.

Usage:  weapon_runtime.py <tokens.txt> [--md]
"""

import math
import os
import re
import struct
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import gmem  # noqa: E402

# vtable VA -> (class name, object stride)
CLASSES = {
    0x820AF548: ("Weapon", 0xC0),
    0x820AF58C: ("Shell", 0x200),
}
NAME_STR_OFF = 0x10  # the name string sits +0x10 into a name record


# ---------------------------------------------------------------- disc side


def read_tokens(path):
    """({class: {ID: {field: value}}}, {(class, ID): [field...]}) from the dump.

    A handful of IDs are declared by more than one `.tbl`. Where two such
    declarations disagree on a field, that field is *ambiguous on disc* — the
    runtime holds whichever definition won, so it must not be scored as a
    contradiction. Those are collected separately, not silently merged.
    """
    seen, ambiguous = {}, {}
    cur = None
    for line in open(path):
        if line.startswith("REC "):
            _, _, _, cls, rid = line.split()
            cur = (cls, rid)
            seen.setdefault(cur, [])
            seen[cur].append({})
        elif line.startswith("F ") and cur is not None:
            _, k, v = line.rstrip("\n").split(" ", 2)
            seen[cur][-1][k] = v

    out = {}
    for (cls, rid), defs in seen.items():
        merged = {}
        for d in defs:
            merged.update(d)
        if len(defs) > 1:
            bad = {
                k
                for k in {k for d in defs for k in d}
                if len({d.get(k) for d in defs}) > 1
            }
            if bad:
                ambiguous[(cls, rid)] = sorted(bad)
                for k in bad:
                    merged.pop(k, None)
        out.setdefault(cls, {})[rid] = merged
    return out, ambiguous


# ------------------------------------------------------------- runtime side


def scan_vtable(f, size, vt):
    pat = struct.pack(">I", vt)
    hits = []
    for a, b in gmem.extents(f.fileno(), size):
        f.seek(a)
        blob = f.read(b - a)
        for m in re.finditer(re.escape(pat), blob):
            off = a + m.start()
            if off % 4 == 0:
                hits.append(off)
    return sorted(set(hits))


def runtime_objects(f, size):
    """{class: {ID: raw_bytes}} plus {class: [(va, ID)]} in array order."""
    objs, order = {}, {}
    for vt, (cls, stride) in CLASSES.items():
        objs[cls], order[cls] = {}, []
        for off in scan_vtable(f, size, vt):
            f.seek(off)
            raw = f.read(stride)
            (nameptr,) = struct.unpack_from(">I", raw, 4)
            try:
                f.seek(gmem.va_to_off(nameptr + NAME_STR_OFF))
                name = f.read(64).split(b"\0")[0].decode("latin-1")
            except (ValueError, UnicodeDecodeError):
                continue
            objs[cls][name] = raw
            order[cls].append((gmem.primary_va(off), name))
    return objs, order


# ------------------------------------------------------------------- solver

ENC = {
    "f32": lambda raw, o, ch: struct.unpack_from(">f", raw, o)[0],
    "u32": lambda raw, o, ch: float(struct.unpack_from(">I", raw, o)[0]),
    "deg": lambda raw, o, ch: math.degrees(struct.unpack_from(">f", raw, o)[0]),
    # A charging weapon drains its magazine continuously, so the two counter
    # fields are float32 on `IsCharging = Yes` records and int32 on the rest.
    # Discovered from the runtime: `IsCharging` partitions the 8 float-encoded
    # records exactly (see docs/re/weapon-struct-runtime.md).
    "cnt": lambda raw, o, ch: (
        struct.unpack_from(">f", raw, o)[0] if ch else float(struct.unpack_from(">I", raw, o)[0])
    ),
}


def near(a, b):
    if a == b:
        return True
    return abs(a - b) <= 2e-4 * max(abs(a), abs(b), 1e-6)


def solve(disc_cls, run_cls, stride, charging):
    """{field: (off, enc, n_agree, [mismatch...])}"""
    numeric = {}
    for rid, fields in disc_cls.items():
        if rid not in run_cls:
            continue
        for k, v in fields.items():
            try:
                numeric.setdefault(k, {})[rid] = float(v.rstrip("fF"))
            except ValueError:
                pass

    solved = {}
    for field, samples in numeric.items():
        if len(samples) < 2:
            continue
        cands = []
        for off in range(0, stride - 3, 4):
            for enc, fn in ENC.items():
                agree, bad = 0, []
                for rid, want in samples.items():
                    got = fn(run_cls[rid], off, rid in charging)
                    if math.isfinite(got) and near(got, want):
                        agree += 1
                    else:
                        bad.append((rid, want, got))
                if agree >= 2:
                    cands.append((len(bad), -agree, off, enc, agree, bad))
        if not cands:
            continue
        cands.sort()
        nbad, _, off, enc, agree, bad = cands[0]
        # A binding that contradicts a large slice of the evidence is not a
        # binding at all -- it is a coincidence on the agreeing subset.
        if len(bad) > 0.2 * len(samples):
            continue
        # Discriminating power: a field whose on-disc samples are all the same
        # value matches any offset holding that constant, so such a binding is
        # a coincidence waiting to happen. Count the distinct values that the
        # *agreeing* records pin down.
        distinct = len({round(v, 6) for rid, v in samples.items() if rid not in {b[0] for b in bad}})
        solved[field] = (off, enc, agree, bad, distinct)

    # Two fields cannot share one byte offset. Where the solver lands two on the
    # same (offset, enc), keep the better-evidenced one and drop the other.
    best_at = {}
    for field, (off, enc, agree, bad, distinct) in solved.items():
        key = (off, enc)
        score = (distinct, agree, -len(bad))
        if key not in best_at or score > best_at[key][1]:
            best_at[key] = (field, score)
    winners = {f for f, _ in best_at.values()}
    return {f: v for f, v in solved.items() if f in winners}, numeric


def report(cls, disc_cls, run_cls, stride, out, charging):
    solved, numeric = solve(disc_cls, run_cls, stride, charging)
    p = out.append
    p(f"\n### `{cls}` — runtime struct ({stride:#x} bytes/object)\n")
    p("Confidence: ✅ = ≥10 disc records agree on ≥3 distinct values, none contradict.")
    p("🟡 = consistent but thin evidence. ⚠️ = the runtime contradicts the disc "
      "(listed below the table).\n")
    p("| offset | enc | field | agree | distinct values | contradict | conf |")
    p("|--------|-----|-------|------:|----------------:|-----------:|------|")
    contradictions = []
    for field, (off, enc, agree, bad, distinct) in sorted(solved.items(), key=lambda kv: kv[1][0]):
        if bad:
            conf = "⚠️"
            contradictions.append((field, off, enc, bad))
        elif agree >= 10 and distinct >= 3:
            conf = "✅"
        else:
            conf = "🟡"
        p(f"| `+{off:#05x}` | {enc} | `{field}` | {agree} | {distinct} | {len(bad)} | {conf} |")

    unsolved = sorted(set(numeric) - set(solved))
    if unsolved:
        p(f"\nNumeric fields with no consistent offset (unsolved): "
          f"{', '.join('`'+u+'`' for u in unsolved)}")

    for field, off, enc, bad in contradictions:
        p(f"\n**`{field}` (`+{off:#05x}`) — runtime disagrees with disc on "
          f"{len(bad)} record(s):**\n")
        for rid, want, got in sorted(bad)[:24]:
            p(f"- `{rid}`: disc `{want:g}` → runtime `{got:g}`")

    p(f"\n#### `{cls}` values the disc defaults, read from the live objects\n")
    for field, (off, enc, agree, bad, distinct) in sorted(solved.items(), key=lambda kv: kv[1][0]):
        if bad or agree < 10 or distinct < 3:
            continue  # only report from ✅ bindings
        missing = [
            (rid, ENC[enc](run_cls[rid], off, rid in charging))
            for rid in sorted(disc_cls)
            if rid in run_cls and field not in disc_cls[rid]
        ]
        if not missing:
            continue
        vals = sorted({round(v, 6) for _, v in missing})
        p(f"\n**`{field}`** (`+{off:#05x}`, {enc}) — defaulted by {len(missing)} of "
          f"{len(run_cls)} records")
        if len(vals) == 1:
            p(f"\n> all = **{vals[0]:g}**")
        else:
            p("")
            for rid, v in missing:
                p(f"- `{rid}` = **{v:g}**")
    return solved


def main():
    tokens = sys.argv[1] if len(sys.argv) > 1 else "/tmp/wep_tokens.txt"
    disc, ambiguous = read_tokens(tokens)
    path = gmem.mem_path()
    size = os.path.getsize(path)
    with open(path, "rb") as f:
        objs, order = runtime_objects(f, size)

    # `IsCharging = Yes` switches the two counter fields to float32; the Shell
    # inherits the flag from its Weapon (same ID suffix).
    charging = {
        rid for rid, f in disc.get("Weapon", {}).items() if f.get("IsCharging") == "Yes"
    }
    charging |= {"Shell" + rid[len("Weapon"):] for rid in set(charging)}

    out = []
    for cls, stride in [(c, s) for c, s in CLASSES.values()]:
        d, r = disc.get(cls, {}), objs.get(cls, {})
        matched = set(d) & set(r)
        out.append(f"\n<!-- {cls}: {len(r)} runtime objects, {len(d)} disc records, "
                   f"{len(matched)} matched by ID -->")
        report(cls, d, r, stride, out, charging)
    if "--csv" in sys.argv:
        import csv

        w = csv.writer(sys.stdout)
        w.writerow(["class", "id", "field", "offset", "enc", "value", "source", "conf"])
        for cls, stride in [(c, st) for c, st in CLASSES.values()]:
            d, r = disc.get(cls, {}), objs.get(cls, {})
            solved, _ = solve(d, r, stride, charging)
            for field, (off, enc, agree, bad, distinct) in sorted(
                solved.items(), key=lambda kv: kv[1][0]
            ):
                conf = "confirmed" if (not bad and agree >= 10 and distinct >= 3) else "tentative"
                for rid in sorted(r):
                    val = ENC[enc](r[rid], off, rid in charging)
                    src = "disc" if field in d.get(rid, {}) else "defaulted-on-disc"
                    w.writerow([cls, rid, field, f"{off:#05x}", enc, f"{val:g}", src, conf])
        return

    if ambiguous:
        out.append("\n### Disc records declared twice, with conflicting values\n")
        out.append("These fields are ambiguous *on disc*; the runtime shows which "
                   "declaration won. They are excluded from the scoring above.\n")
        for (cls, rid), fields in sorted(ambiguous.items()):
            out.append(f"- `{cls}` `{rid}`: {', '.join('`'+f+'`' for f in fields)}")
    print("\n".join(out))


if __name__ == "__main__":
    main()