#!/usr/bin/env python3
"""Find the afterburner by what it COSTS, not by what it speeds up.

`ab_probe.py` held A/B/X/LB at full throttle and none of them moved the speed —
so either the burner is a different input or it does not raise the speed target
(the definition has no `AB_*Velocity`, only `AB_AV_*` turn caps and
`AB_ConsumeShield{,_Begin}` 10 / 50). This looks for the cost instead: sample a
window of the player object while each candidate is held, and report any float
that FALLS during the hold and not before it — `AB_ConsumeShield_Begin 50` should
be a visible step.

Usage: ab_state_probe.py [hold_s]
"""
import os, struct, subprocess, sys, time
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import speed_law

WIN_LO, WIN_HI = 0x100, 0x220          # object window, relative to the position triple

def pad(*a):
    subprocess.run(["vgamepad", *a], capture_output=True)

def snap(fd, off):
    b = os.pread(fd, WIN_HI - WIN_LO, off + WIN_LO)
    return [struct.unpack_from(">f", b, i)[0] for i in range(0, len(b) - 3, 4)]

def main():
    hold = float(sys.argv[1]) if len(sys.argv) > 1 else 4.0
    w, off, nm = speed_law.find_player()
    if not w:
        sys.exit("player entity not found after retries")
    print(f"# locked on {nm}; window {WIN_LO:#x}..{WIN_HI:#x}")
    pad("trig", "RT", "1.0")
    time.sleep(1.5)
    for btn in ("LS", "RS", "A", "B", "X", "LB"):
        before = snap(w.fd, off)
        pad("press", btn)
        time.sleep(hold)
        during = snap(w.fd, off)
        pad("release", btn)
        time.sleep(1.0)
        after = snap(w.fd, off)
        drops = []
        for i, (b, d, a) in enumerate(zip(before, during, after)):
            # a float that fell while held and did not fall further after release
            if b - d > 1.0 and (a - d) > -1.0 and abs(b) < 1e6:
                drops.append((WIN_LO + i * 4, round(b, 1), round(d, 1), round(a, 1)))
        tag = ", ".join(f"{o:#05x}: {b}->{d}->{a}" for o, b, d, a in drops[:4]) or "nothing fell"
        print(f"# {btn:<3} {tag}")
    pad("trig", "RT", "0.0"); pad("reset")

if __name__ == "__main__":
    main()