#!/usr/bin/env python3
"""Drive a scripted input sequence in flight while sampling guest RAM, so the
player craft's live transform can be *derived* rather than guessed.

The point is correlation: each sample is tagged with the stick/trigger state
that produced it, so afterwards we can ask "which floats move only while the
yaw axis is deflected?" and "which triple integrates to the velocity?".

Writes a .npz-ish flat binary: a header of (va, offset, name) per instance,
then per frame a timestamp, the input vector, and every instance's raw window.

Usage: flight_probe.py <out.bin> [seconds]
"""
import os
import struct
import sys
import time

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

FIFO = "/tmp/sylph-vgamepad.fifo"


class Pad:
    """Talk to the vgamepad server directly over its FIFO.

    The `vgamepad` CLI spawns a process per command (~20 ms); a control loop
    cannot afford that, and `tap`/`hold` additionally sleep *inside* the server.
    Writing lines to the FIFO ourselves keeps a tick under a millisecond.
    """

    def __init__(self):
        self.f = open(FIFO, "w", buffering=1)
        self.state = {"LX": 0.0, "LY": 0.0, "RX": 0.0, "RY": 0.0, "LT": 0.0, "RT": 0.0}

    def axis(self, name, v):
        self.state[name] = v
        self.f.write(f"axis {name} {v:.3f}\n")

    def trig(self, name, v):
        self.state[name] = v
        self.f.write(f"trig {name} {v:.3f}\n")

    def press(self, b):
        self.f.write(f"press {b}\n")

    def release(self, b):
        self.f.write(f"release {b}\n")

    def reset(self):
        self.f.write("reset\n")
        for k in self.state:
            self.state[k] = 0.0

    def vector(self):
        return [self.state[k] for k in ("LX", "LY", "RX", "RY", "LT", "RT")]


# (duration_s, description, action)
SCRIPT = [
    (3.0, "idle", lambda p: p.reset()),
    (4.0, "pitch-up", lambda p: p.axis("LY", -0.9)),
    (2.0, "idle2", lambda p: p.reset()),
    (4.0, "yaw-right", lambda p: p.axis("LX", 0.9)),
    (2.0, "idle3", lambda p: p.reset()),
    (4.0, "yaw-left", lambda p: p.axis("LX", -0.9)),
    (2.0, "idle4", lambda p: p.reset()),
    (5.0, "boost", lambda p: p.trig("RT", 1.0)),
    (3.0, "idle5", lambda p: p.reset()),
]


def main():
    out = sys.argv[1]
    w = gworld.World()
    inst = w.refresh()
    print(f"# {len(inst)} instances", flush=True)
    if not inst:
        sys.exit("no instances — not in a mission?")

    pad = Pad()
    hz = 10.0
    with open(out, "wb") as f:
        f.write(b"SYLPHPRB")
        f.write(struct.pack("<III", len(inst), gworld.WINDOW, 6))
        for va, off, nm in inst:
            f.write(struct.pack("<IQ", va, off) + nm.encode()[:63].ljust(64, b"\0"))
        t_start = time.time()
        for dur, label, act in SCRIPT:
            act(pad)
            print(f"# {label} ({dur}s)", flush=True)
            n = int(dur * hz)
            for _ in range(n):
                t = time.time()
                f.write(struct.pack("<d", t - t_start))
                f.write(struct.pack("<6f", *pad.vector()))
                f.write(struct.pack("<32s", label.encode()[:32]))
                for va, off, nm in inst:
                    f.write(w.read_off(off, gworld.WINDOW).ljust(gworld.WINDOW, b"\0"))
                time.sleep(max(0, 1.0 / hz - (time.time() - t)))
        pad.reset()
    print(f"# wrote {out} ({os.path.getsize(out)/1e6:.1f} MB)", flush=True)


if __name__ == "__main__":
    main()