#!/usr/bin/env python3
"""firstop.py — catch the plugin INSIDE its audio callback by repeatedly
SIGSTOP-ing the yabridge host and reading only the ctx-referenced arrays.

Between callbacks the FIR work buffer (ctx+0x540668) is reset to the complex
identity (1,0)x2049. A snapshot where FIR != identity means we stopped after
the mask->FIR build stage; those snapshots are saved with the full pipeline
state (bands/scratch/f6f8/track/R curves).

Usage: python3 scripts/firstop.py [rpp] [nattempts]
"""
import os
import signal
import struct
import subprocess
import sys
import time

import numpy as np

SNAPDIR = '/tmp/opencode/firstop'
SLOTS = [0x540548, 0x540550, 0x540598, 0x540628, 0x540668, 0x540678, 0x540688,
         0x540698, 0x5406a8, 0x5406b8, 0x5406c8, 0x5406d8, 0x5406e8, 0x5406f8,
         0x540708, 0x540718, 0x540728, 0x540738, 0x540748, 0x540758,
         0x540768, 0x540778, 0x540788, 0x540798, 0x5407a8, 0x5407b8,
         0x5407c8, 0x5407d8, 0x5407e8, 0x5407f8, 0x540808, 0x540818,
         0x540828, 0x540838, 0x540848]
NARR = 8194


def find_host():
    import glob
    for p in glob.glob('/proc/[0-9]*'):
        pid = int(os.path.basename(p))
        try:
            cmd = open(f'/proc/{pid}/cmdline', 'rb').read().replace(b'\0', b' ').decode('utf8', 'replace')
            maps = open(f'/proc/{pid}/maps').read()
        except Exception:
            continue
        if 'soothe2' in maps and 'reaper' not in cmd:
            return pid
    return None


def main():
    rpp = sys.argv[1] if len(sys.argv) > 1 else '/home/m/soothe-bt/dual_b1q_0.5.rpp'
    nattempts = int(sys.argv[2]) if len(sys.argv) > 2 else 120
    os.makedirs(SNAPDIR, exist_ok=True)
    subprocess.run('pkill -9 -x reaser 2>/dev/null; pkill -9 -x reaper 2>/dev/null; '
                   "pkill -9 -f '[y]abridge' 2>/dev/null; sleep 1", shell=True)
    proc = subprocess.Popen(
        ['/usr/bin/reaper', '-nosplash', '-ignoreerrors', rpp,
         '/home/m/re-tools/play_loop.lua'],
        stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT)
    t0 = time.time()
    host = None
    while time.time() - t0 < 60 and not host:
        host = find_host()
        time.sleep(0.2)
    if not host:
        print('NO HOST')
        return 1
    print('host', host, flush=True)
    time.sleep(12)

    fd = os.open(f'/proc/{host}/mem', os.O_RDONLY)

    def rd(addr, n):
        try:
            return os.pread(fd, n, addr)
        except OSError:
            return None

    def get_ctx():
        # cheap probe: sens scalar must be >100 at known offset
        for base in (CTX.get('base'),):
            pass
        return CTX.get('base')

    # locate ctx once (while running): marker scan over heap only is heavy;
    # reuse known-good address from prior sessions, validate via sens scalar.
    CTX = {}
    ctx = None
    b = rd(0x2370040 + 0x540870, 4)
    if b and struct.unpack('<f', b)[0] > 100:
        ctx = 0x2370040
    else:
        print('known ctx invalid; full scan needed')
        return 1
    print('ctx', hex(ctx), flush=True)

    hits = 0
    saved = []
    rng = np.random.default_rng(7)
    for k in range(nattempts):
        os.kill(host, signal.SIGSTOP)
        try:
            fir_p = struct.unpack('<Q', rd(ctx + 0x540668, 8))[0]
            fb = rd(fir_p, 64 * 4)
            ident = False
            if fb:
                arr = np.frombuffer(fb[:256], dtype='<f4')
                ident = bool(np.all(np.abs(arr[0::2] - 1.0) < 1e-6))
                ident = ident and bool(np.all(arr[1::2] == 0))
            if not ident:
                hits += 1
                store = {}
                for off in SLOTS:
                    pb = rd(ctx + off, 8)
                    if not pb:
                        continue
                    p = struct.unpack('<Q', pb)[0]
                    if p < 0x10000:
                        continue
                    ab = rd(p, NARR * 4)
                    if not ab:
                        continue
                    store[hex(off)] = np.frombuffer(ab, dtype='<f4').astype(np.float32)
                fn = f'{SNAPDIR}/hit{k:03d}.npz'
                np.savez_compressed(fn, **store)
                saved.append(fn)
                sc = rd(ctx + 0x2404dc, 4)
                print(f'[{k}] HIT fir!=identity -> {fn} '
                      f'(fir[0..5]={np.round(store["0x540668"][:6],4)})', flush=True)
        finally:
            os.kill(host, signal.SIGCONT)
        time.sleep(float(rng.uniform(0.02, 0.09)))
    os.close(fd)
    if proc.poll() is None:
        proc.kill()
    print(f'done: {hits} hits / {nattempts} attempts -> {SNAPDIR}')
    return 0


if __name__ == '__main__':
    sys.exit(main())