#!/usr/bin/env python3
"""rtctx.py — P1.5: live-capture of level-tracker coefficients + mask scalars.

Pipeline:
  1. spawn reaper offline render of render_long.rpp (background),
  2. wait for yabridge-host with soothe2 mapped,
  3. chunked full-heap snapshot to /tmp/snap_all.bin (like rtsnap_fast.py),
  4. locate the registry {u64 count,u64 ptr} run (NOTES_CAPTURE),
  5. derive ctx base = (field addr holding window ptr) - 0x540658,
  6. dump the P1.5 targets: level-tracker A[] (0x4c0528/0x3c0510/0x2c04f8)
     and mask scalars (0x540870/874/87c/888/88c) + BandConfig (ctx+0x188).

Usage: rtctx.py [--keep]
"""
import subprocess, os, glob, time, struct, sys, json

KEEP = '--keep' in sys.argv[1:]
SNAP = '/tmp/snap_all.bin'
IDX = '/tmp/snap_all.idx'
RPP = '/home/m/soothe-bt/render_long.rpp'


def mem_open(pid):
    return os.open(f'/proc/{pid}/mem', os.O_RDONLY)


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):
    fd = mem_open(host)
    out = open(SNAP, 'wb')
    idx = open(IDX, 'wb')
    t0 = time.time(); 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)
            idx.write(struct.pack('<QQ', a, len(d)))
            nreg += 1; nbytes += len(d)
            a += n
    out.close(); idx.close()
    return nreg, nbytes


def parse_snap():
    data = open(SNAP, 'rb').read()
    i = 0; regs = []
    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_registry(regs):
    """Look for {u64 count,u64 ptr} registry run (NOTES_CAPTURE):
    a run of >=4 pairs at stride 0x10, first count==8193 (window registry),
    every count in {8193,2049,2048,32768,16384,1024} and each ptr readable."""
    okc = {8193, 2049, 2048, 32768, 16384, 1024, 4096}
    min_run = 4
    hits = []
    for lo, body in regs:
        if len(body) < 0x2000:
            continue
        for off in range(0, len(body) - 16 * min_run, 8):
            good = True
            for k in range(min_run):
                if off + 16 + 8 > len(body):
                    good = False; break
                c = struct.unpack_from('<Q', body, off + k * 16)[0]
                if c not in okc:
                    good = False; break
                if off + k * 16 + 8 > len(body):
                    good = False; break
                p = struct.unpack_from('<Q', body, off + k * 16 + 8)[0]
                if readabs(regs, p, 4) is None:
                    good = False; break
            if good:
                hits.append(lo + off)
    # dedup: drop candidates too close (same run)
    hits = sorted(set(hits))
    dedup = []
    for h in hits:
        if not dedup or h - dedup[-1] > 0x40:
            dedup.append(h)
    return dedup


def find_ctx(regs, winptr):
    """Find the DSP ctx: field at base+0x24 == 40000.0f (0x472c4400),
    and field at base+0x540658 holds the window ptr (winptr)."""
    sig = struct.pack('<I', 0x472c4400)
    for lo, body in regs:
        if len(body) < 0x541000:
            continue
        i = 0
        while True:
            i = body.find(sig, i)
            if i < 0:
                break
            base = lo + i - 0x24
            if base >= lo:
                # verify window field
                b = readabs(regs, base + 0x540658, 8)
                if b and struct.unpack('<Q', b)[0] == winptr:
                    return base
            i += 1
    return None


def dump_targets(regs, ctx):
    out = {}
    def f32(addr):
        b = readabs(regs, addr, 4)
        return struct.unpack('<f', b)[0] if b else None
    def f64arr(addr, n):
        b = readabs(regs, addr, 8 * n)
        return list(struct.unpack('<%dd' % n, b)) if b and len(b) >= 8 * n else None

    out['scalars'] = {
        '0x24': f32(ctx + 0x24),         # 40000.0 (sr?)
        '0x540870': f32(ctx + 0x540870),  # sens/gain
        '0x540874': f32(ctx + 0x540874),  # level-dependent weight
        '0x54087c': f32(ctx + 0x54087c),  # band weight / mix
        '0x540888': f32(ctx + 0x540888),  # attack coeff = 10^(att/20)
        '0x54088c': f32(ctx + 0x54088c),  # release coeff
    }
    # BandConfig at ctx+0x188 (A=+0 B=+4 gamma=+0c flag=+10)
    out['bandconfig0'] = {
        'A': f32(ctx + 0x188),
        'B': f32(ctx + 0x188 + 0x04),
        'gamma': f32(ctx + 0x188 + 0x0c),
        'flag': readabs(regs, ctx + 0x188 + 0x10, 1),
    }
    # level-tracker coefficient arrays (first-order IIR A[])
    for name, off, cnt in [('0x4c0528', 0x4c0528, 32),
                           ('0x3c0510', 0x3c0510, 32),
                           ('0x2c04f8', 0x2c04f8, 32)]:
        a = f64arr(ctx + off, cnt)
        if a is not None:
            out[name] = a
    return out


def main():
    os.system("pkill -9 -x reaper 2>/dev/null; pkill -9 -f '[y]abridge' 2>/dev/null; sleep 1")
    proc = subprocess.Popen(['/usr/bin/reaper', '-nosplash', '-ignoreerrors', '-renderproject', RPP],
                            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)
    time.sleep(6)
    nreg, nbytes = snapshot(host)
    print('snapshot %d regs %.1f MB' % (nreg, nbytes / 1e6))
    if proc.poll() is None: proc.kill()

    regs = parse_snap()
    hits = find_registry(regs)
    print('registry candidates:', len(hits))
    # choose the window-registry: first pair count==8193 (identity/window).
    rb = None
    for h in hits:
        b = readabs(regs, h, 8)
        if b:
            c0 = struct.unpack('<Q', b)[0]
            if c0 == 8193:
                rb = h
                break
    if rb is None:
        print('NO WINDOW-REGISTRY (first cnt==8193)'); return 1
    win = struct.unpack_from('<Q', readabs(regs, rb + 8, 8))[0]
    n00 = struct.unpack_from('<Q', readabs(regs, rb, 8))[0]
    print('registry base=0x%x [00].cnt=%d [01].ptr(window)=0x%x' % (rb, n00, win))
    ctx = find_ctx(regs, win)
    if not ctx:
        print('NO CTX'); return 1
    print('CTX base = 0x%x' % ctx)
    dump = dump_targets(regs, ctx)
    print(json.dumps(dump, indent=1, default=str))
    with open('/tmp/rtctx_dump.json', 'w') as f:
        json.dump(dump, f, indent=1, default=str)
    print('saved /tmp/rtctx_dump.json')
    return 0


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