#!/usr/bin/env python3
"""dualtrace.py — live capture of the DSP ctx during realtime playback of
dual_b1q_0.5.rpp (two tones 500+2000 Hz, one band @500 q=0.5).

Method (NOTES_CAPTURE.md 2026-08-20c): reaper + play_loop.lua keeps the audio
callback alive; chunked /proc/pid/mem snapshot of the yabridge host; find ctx
by marker +0x24 == 48000.0f; dump scalars + pointer table + all ~2049-float
arrays, flag those with energy at tone bins 43/171.

Usage: python3 scripts/dualtrace.py [rpp]
"""
import json
import os
import struct
import subprocess
import sys
import time

import numpy as np

SNAP1 = '/tmp/opencode/dualtrace_s1.bin'
SNAP2 = '/tmp/opencode/dualtrace_s2.bin'


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 snapshot(host, path):
    fd = os.open(f'/proc/{host}/mem', os.O_RDONLY)
    out = open(path, 'wb')
    nreg = nbytes = 0
    for line in open(f'/proc/{host}/maps').read().splitlines():
        p = line.split()
        if len(p) < 2 or 'r' not in p[1]:
            continue
        lo, hi = (int(x, 16) for x in p[0].split('-'))
        a = lo
        while a < hi:
            n = min(hi - a, 8 * 1024 * 1024)
            try:
                d = os.pread(fd, n, a)
            except Exception:
                a += n
                continue
            if d:
                out.write(struct.pack('<QQ', a, len(d)))
                out.write(d)
                nreg += 1
                nbytes += len(d)
            a += n
    out.close()
    os.close(fd)
    return nreg, nbytes


def parse_snap(path):
    data = open(path, 'rb').read()
    regs = []
    i = 0
    while i + 16 <= len(data):
        lo, sz = struct.unpack_from('<QQ', data, i)
        regs.append((lo, data[i + 16:i + 16 + sz]))
        i += 16 + sz
    return regs


def readabs(regs, addr, n):
    for lo, body in regs:
        if lo <= addr < lo + len(body) and addr - lo + n <= len(body):
            return body[addr - lo:addr - lo + n]
    return None


def find_ctx(regs):
    sig = struct.pack('<I', 0x473b8000)  # 48000.0f
    cands = []
    for lo, body in regs:
        j = 0
        while True:
            j = body.find(sig, j)
            if j < 0:
                break
            base = lo + j - 0x24
            b = readabs(regs, base + 0x540870, 4)
            if b and struct.unpack('<f', b)[0] > 100:  # sens scalar ~441
                cands.append(base)
            j += 1
    return sorted(set(cands))


def dump_ctx(regs, ctx, tag, store):
    """Scalars + pointer table 0x540600..0x540900 + dereferenced arrays."""
    def f32(addr):
        b = readabs(regs, addr, 4)
        return struct.unpack('<f', b)[0] if b else None

    scal = {hex(a): f32(ctx + a) for a in range(0x540860, 0x5408a8, 4)}
    store[f'{tag}_scalars'] = scal

    arrays = {}
    ptrinfo = []
    for off in range(0x540600, 0x540900, 8):
        b = readabs(regs, ctx + off, 8)
        if not b:
            continue
        ptr = struct.unpack('<Q', b)[0]
        if ptr < 0x10000 or ptr > 0x7fffffffffff:
            continue
        body = readabs(regs, ptr, 2049 * 4 + 64)
        if body is None or len(body) < 2049 * 4:
            continue
        arr = np.frombuffer(body[:2049 * 4], dtype='<f4').astype(np.float64)
        if not np.isfinite(arr).all():
            continue
        arrays[off] = arr
        ptrinfo.append((hex(ctx + off), hex(ptr)))

    # classify: energy at tone bins relative to local median
    flagged = {}
    med = None
    for off, arr in arrays.items():
        m = float(np.median(arr))
        if m <= 0 or not np.isfinite(m) or float(np.max(np.abs(arr))) > 1e6:
            continue
        r43 = arr[43] / m
        r171 = arr[171] / m
        if r43 > 3 or r171 > 3 or (arr.max() / max(m, 1e-30)) > 5:
            flagged[off] = dict(ratio43=r43, ratio171=r171,
                                vmin=float(arr.min()), vmax=float(arr.max()))
            store[f'{tag}_arr_{hex(off)}'] = arr
    store[f'{tag}_ptrinfo'] = ptrinfo
    store[f'{tag}_flagged'] = {hex(k): v for k, v in flagged.items()}
    return flagged


def main():
    rpp = sys.argv[1] if len(sys.argv) > 1 else '/home/m/soothe-bt/dual_b1q_0.5.rpp'
    os.makedirs('/tmp/opencode', 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, 'at %.1fs' % (time.time() - t0), flush=True)
    time.sleep(12)  # init + several looped passes -> steady state

    store = {}
    try:
        n1, b1 = snapshot(host, SNAP1)
        print(f'snap1: {n1} regs {b1/1e6:.0f}MB', flush=True)
        time.sleep(6)
        n2, b2 = snapshot(host, SNAP2)
        print(f'snap2: {n2} regs {b2/1e6:.0f}MB', flush=True)

        for tag, path in (('s1', SNAP1), ('s2', SNAP2)):
            regs = parse_snap(path)
            cands = find_ctx(regs)
            print(tag, 'ctx candidates:', [('0x%x' % c) for c in cands][:5], flush=True)
            if cands:
                fl = dump_ctx(regs, cands[0], tag, store)
                for k, v in list(fl.items())[:10]:
                    print(' ', k, {kk: round(vv, 2) if isinstance(vv, float) else vv
                                   for kk, vv in v.items()}, flush=True)
    finally:
        if proc.poll() is None:
            proc.kill()

    np.savez_compressed('/tmp/opencode/dualtrace.npz',
                        **{k: v for k, v in store.items()
                           if isinstance(v, np.ndarray)})
    json.dump({k: v for k, v in store.items() if not isinstance(v, np.ndarray)},
              open('/tmp/opencode/dualtrace.json', 'w'), indent=1)
    print('\nsaved /tmp/opencode/dualtrace.npz|.json')
    return 0


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