#!/usr/bin/env python3
"""Per-call-site RNG ledger from a lane-Z trace with the draw-site detours enabled.

The boundary ledger (tools/rng_ledger_report.py) says how many words a turn spent and in which
phase. This says which CALL SITE spent them: the seven generator entry points are detoured and each
call records __builtin_return_address(0) -- the game instruction after its own `call` -- with the
word cost taken from `left` before and after.

Two corrections are applied here rather than in the shim, because both are bookkeeping:

  * rows whose return address falls INSIDE another entry point's body are the helper's own internal
    draw (Chance calls NextFloat; IntRangeBell calls NextInt twice). Those words are already counted
    against the helper's row, so including them double-counts.
  * rows tagged `strategic: false` came from a different Mars::RNG instance (the StrategyClient's at
    +0x134, the tactical CombatSim's, a map-generation temporary). They do not touch the save's
    strategic generator and must not be added to the turn's total.

The sum of the remaining rows must equal the bracket total the boundary ledger measured
independently. A shortfall is an unattributed word.

  uv run python3 tools/rng_site_report.py <trace.jsonl>
"""
import json, sys, bisect, os

# entry-point bodies (RVAs), for the helper-internal test
BODIES = {"Chance": (0x4E6DD0, 0x4E6E2E), "FloatRange": (0x7D8A0, 0x7D8C9),
          "IntRangeBell": (0x4E6D80, 0x4E6DC6), "GaussianRange": (0x4E6E30, 0x4E6FCA)}

def internal_of(rva):
    for name, (lo, hi) in BODIES.items():
        if lo <= rva <= hi:
            return name
    return None

fns, addrs = {}, []
fp = os.path.expanduser("~/sots-re/dumps/functions.json")
if os.path.exists(fp):
    fns = json.load(open(fp))
    addrs = sorted(int(k, 16) for k in fns)

def owner(va):
    if not addrs:
        return ""
    i = bisect.bisect_right(addrs, va) - 1
    if i < 0:
        return ""
    a = addrs[i]
    name, _ = fns[f"0x{a:08x}"]
    return f"{name}+0x{va - a:x}"

for line in open(sys.argv[1]):
    d = json.loads(line)
    if "hook" not in d or not d["hook"].endswith("Autosave"):
        continue
    a = {x.get("n"): x for x in d.get("args", []) if x.get("n")}
    if a.get("end_turn", {}).get("v") is not False:
        continue
    rows = [r["v"] for r in a.get("draw_sites", {}).get("v", [])]
    print(f"turn: strategic {a.get('draw_site_words',{}).get('v')} words / "
          f"{a.get('draw_site_calls',{}).get('v')} calls; "
          f"other generators {a.get('draw_site_words_other_rng',{}).get('v')} words / "
          f"{a.get('draw_site_calls_other_rng',{}).get('v')} calls; "
          f"overflow {a.get('draw_site_overflow',{}).get('v')}")
    net = 0
    for v in sorted(rows, key=lambda r: (not r["strategic"]["v"], r["ret_rva"]["v"])):
        rva = v["ret_rva"]["v"]
        ins = internal_of(rva)
        strat = v["strategic"]["v"]
        counted = strat and not ins
        if counted:
            net += v["words"]["v"]
        call = rva - 5  # every entry point here is reached by a 5-byte call rel32
        tag = "STRAT" if strat else "other"
        note = f"   [internal to {ins}: already counted on that row]" if ins else ""
        print(f"  {'*' if counted else ' '} {tag} {v['entry']['v']:<13} "
              f"call 0x{call+0x400000:08x} {owner(call + 0x400000):<34} "
              f"calls={v['calls']['v']:>3} words={v['words']['v']:>3}{note}")
    print(f"  => attributed strategic words: {net}")