#!/usr/bin/env python3
"""step7_capture.py — Step 7: live capture of level-path BandConfig A/B/gamma per RPP.

Method (NOTES_CAPTURE 2026-08-20c):
  1. spawn reaper with <rpp> (+ play.lua realtime transport by default),
  2. wait for yabridge-host (soothe2 mapped, not reaper), sleep for init,
  3. chunked full-heap snapshot (8MB pread chunks; large regions EIO otherwise),
  4. fingerprint scan: per-band level_gain pair buffers = 0x400 [level,gain]
     f32 pairs with level[j] == j/1024 EXACTLY (j<1024 -> exact in fp32),
  5. u64 refs to those buffers land at base+0xe0+band*0x18 (stride 0x18)
     -> majority-vote the level-path object base,
  6. decode every qword ptr at base+0x170..0x1a8 as BandConfig:
     A@+0x0 B@+0x4 gamma@+0xc flag@+0x10 shaper@+0x18.. callback@+0x90,
     plus first mask doubles at base+0x4198+band*0x2000.

Usage:
  python3 scripts/step7_capture.py <file.rpp> [--offline] [--json PATH] [--snap PATH]

Default mode is realtime playback (play.lua) so short projects keep the DSP
host alive during capture. --offline uses -renderproject instead.
"""
import subprocess, os, glob, time, struct, sys, json, collections

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PLAY_LUA = os.path.join(REPO, 'play.lua')


def find_host():
    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 = 0; nbytes = 0
    for line in open(f'/proc/{host}/maps').read().splitlines():
        p = line.split()
        if len(p) < 2: continue
        lo, hi = (int(x, 16) for x in p[0].split('-'))
        if 'r' not in p[1]: continue
        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 not d:
                a += n; continue
            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):
    regs = []
    data = open(path, 'rb')
    while True:
        hdr = data.read(16)
        if len(hdr) < 16: break
        lo, sz = struct.unpack('<QQ', hdr)
        body = data.read(sz)
        if len(body) < sz: break
        regs.append((lo, body))
    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


import numpy as np


def find_levelpair_buffers(regs):
    """Buffers where f32[2j]==j/1024 exactly for j=0..1023 (level axis)."""
    expect = (np.arange(1024, dtype=np.float64) / 1024).astype(np.float32)
    anchors = np.nonzero(expect == np.float32(1))[0]  # sanity of construction
    found = []
    for lo, body in regs:
        if len(body) < 8192: continue
        a = np.frombuffer(body[:len(body) // 4 * 4], dtype='<f4')
        # anchor: level[1] == 1/1024 at slot 2
        cand = np.nonzero(a == np.float32(1.0 / 1024.0))[0]
        for i in cand:
            i = int(i)
            if i < 2 or i % 2: continue
            s = i - 2  # slot of level[0]
            if s + 2048 > len(a): continue
            if np.array_equal(a[s:s + 2048:2], expect):
                found.append(lo + 4 * s)
    found = sorted(set(found))
    dedup = []
    for h in found:
        if not dedup or h - dedup[-1] > 0x100:
            dedup.append(h)
    return dedup


def find_object_base(regs, bufs):
    """u64 refs to band buffers sit at base+0xe0+band*0x18."""
    votes = collections.Counter()
    detail = []
    for k, buf in enumerate(bufs):
        pat = struct.pack('<Q', buf)
        for lo, body in regs:
            j = 0
            while True:
                j = body.find(pat, j)
                if j < 0: break
                addr = lo + j
                base = addr - 0xe0 - k * 0x18
                votes[base] += 1
                detail.append((k, addr, base))
                j += 1
    if not votes:
        return None, detail
    base, cnt = votes.most_common(1)[0]
    return (base, cnt) if cnt >= 2 else (None, detail)


def looks_like_bandconfig(regs, p):
    b = readabs(regs, p, 0x98)
    if not b: return False
    A, B = struct.unpack_from('<ff', b, 0)
    g, = struct.unpack_from('<f', b, 0xc)
    if not (-96.0 <= A <= 96.0): return False
    if not (1.0 <= B <= 100000.0): return False
    if not (0.05 <= abs(g) <= 8.0): return False
    return True


def dump_bandconfigs(regs, base):
    def f32(addr):
        v = readabs(regs, addr, 4)
        return round(struct.unpack('<f', v)[0], 6) if v else None
    def u64(addr):
        v = readabs(regs, addr, 8)
        return struct.unpack('<Q', v)[0] if v else None
    out = {}
    for off in range(0x80, 0x400, 8):
        p = u64(base + off)
        if not p or p < 0x10000: continue
        if p & 7 or not looks_like_bandconfig(regs, p): continue
        out['+0x%x' % off] = {
            'ptr': '0x%x' % p,
            'A': f32(p), 'B': f32(p + 4),
            'f08': f32(p + 8), 'gamma': f32(p + 0xc),
            'flag': (readabs(regs, p + 0x10, 1) or b'\xff')[0],
            'raw_f32': [f32(p + x) for x in range(0x14, 0x30, 4)],
            'cb': ('set' if u64(p + 0x90) else 'none'),
        }
    masks = {}
    for band in range(6):
        b = readabs(regs, base + 0x4198 + band * 0x2000, 32)
        if b:
            masks['band%d' % band] = [round(v, 4) for v in struct.unpack('<4d', b)]
    return out, masks


def main():
    args = sys.argv[1:]
    rpp = args[0]
    offline = '--offline' in args
    jpath = None
    spath = '/tmp/opencode/step7_snap.bin'
    if '--json' in args: jpath = args[args.index('--json') + 1]
    if '--snap' in args: spath = args[args.index('--snap') + 1]
    name = os.path.splitext(os.path.basename(rpp))[0]
    if not jpath:
        jpath = f'/tmp/opencode/step7_{name}.json'
    os.makedirs(os.path.dirname(jpath), exist_ok=True)

    subprocess.run('pkill -9 -x reaper 2>/dev/null; pkill -9 -f "[y]abridge" 2>/dev/null; sleep 1', shell=True)
    if offline:
        cmd = ['/usr/bin/reaper', '-nosplash', '-ignoreerrors', '-renderproject', rpp]
    else:
        cmd = ['/usr/bin/reaper', '-nosplash', '-ignoreerrors', rpp, PLAY_LUA]
    t0 = time.time()
    proc = subprocess.Popen(cmd, stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT)
    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 %d at %.1fs' % (host, time.time() - t0))
    time.sleep(8)
    for attempt in range(3):
        try:
            nreg, nbytes = snapshot(host, spath)
        except Exception as e:
            print('snapshot failed:', e); break
        print('snapshot #%d: %d regs %.1f MB' % (attempt, nreg, nbytes / 1e6))
        if nbytes > 50e6: break
        time.sleep(3)
    if proc.poll() is None: proc.kill()

    regs = parse_snap(spath)
    bufs = find_levelpair_buffers(regs)
    print('level-pair buffers:', ['0x%x' % b for b in bufs])
    res = {'rpp': rpp, 'buffers': ['0x%x' % b for b in bufs]}
    if not bufs:
        json.dump(res, open(jpath, 'w'), indent=1); print('saved', jpath); return 2
    # group into bands: cluster addrs with stride ~0x2000
    groups = [[bufs[0]]]
    for b in bufs[1:]:
        if b - groups[-1][-1] <= 0x3000: groups[-1].append(b)
        else: groups.append([b])
    best = None
    for g in groups:
        base, cnt = find_object_base(regs, g)
        print('group n=%d -> base %s (votes=%s)' % (len(g), ('0x%x' % base) if base else None, cnt if isinstance(cnt, int) else '-'))
        if base and (best is None or cnt > best[1]):
            best = (base, cnt)
    if not best:
        json.dump(res, open(jpath, 'w'), indent=1); print('saved', jpath); return 3
    base = best[0]
    cfgs, masks = dump_bandconfigs(regs, base)
    res.update(levelpath_base='0x%x' % base, configs=cfgs, masks_head=masks)
    json.dump(res, open(jpath, 'w'), indent=1)
    print(json.dumps(cfgs, indent=1))
    print('masks head:', json.dumps(masks))
    print('saved', jpath)
    return 0


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