#!/usr/bin/env python3
"""rtobj.py — runtime-capture of soothe2 DSP object heap (v2: periodic heap scan).

Spawns `reaper -nosplash -renderproject <rpp>`, locates the yabridge-host process
with soothe2 mapped, computes modbase, and PERIODICALLY scans ONLY anonymous rw
regions (heaps) for the Soothe2Module vtable pointer, verifying by ctor fields.
On a valid object it dumps twin-state / step-5 window / level weights / mask /
warp / freqaxis / setters into OUTDIR.
"""
import subprocess, os, glob, sys, time, struct, json

VPTR_RVA = 0x24abb80
VPTR_N = 0x60
N_BINS = 342
N_DUMP = 2048


def mem_read(pid, addr, n):
    mem = os.open(f'/proc/{pid}/mem', os.O_RDONLY)
    try:
        return os.pread(mem, n, addr)
    except Exception:
        return None
    finally:
        os.close(mem)


def find_host(proc, wait=120):
    t0 = time.time()
    while time.time() - t0 < wait:
        for p in glob.glob('/proc/[0-9]*'):
            try:
                pid = int(os.path.basename(p))
                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 'yabridge-host' in cmd and 'soothe2' in maps:
                return pid
        time.sleep(0.05)
    return None


def soothe2_base(host):
    try:
        maps = open(f'/proc/{host}/maps').read()
    except Exception:
        return None
    for line in maps.splitlines():
        if 'soothe2' not in line:
            continue
        rlo = int(line.split()[0].split('-')[0], 16)
        hdr = mem_read(host, rlo, 0x1000)
        if hdr and hdr[:2] == b'MZ':
            x = struct.unpack('<I', hdr[0x3c:0x40])[0]
            if hdr[x:x + 4] == b'PE\x00\x00':
                return rlo
    return None


def heap_regions(pid, excl_start=None, excl_end=None):
    out = []
    try:
        for line in open(f'/proc/{pid}/maps').read().splitlines():
            p = line.split()
            if len(p) < 2:
                continue
            path = p[5] if len(p) > 5 else ''
            if 'r' not in p[1]:
                continue
            lo, hi = (int(x, 16) for x in p[0].split('-'))
            # skip the image file-back regions of the plugin itself
            if excl_start is not None and excl_start <= lo < excl_end:
                continue
            if 'soothe2' in line:
                continue
            out.append((lo, hi))
    except Exception:
        pass
    return out


def read_vslots(host, base):
    """vtable pointer values = addresses of the .data vtable-slots themselves
    (object's first qword points here), range [vstart, vstart+VPTR_N)."""
    vstart = base + VPTR_RVA
    slots = [vstart + off for off in range(0, VPTR_N, 8)]
    return slots


def scan_objects(host, slots, excl_start=None, excl_end=None):
    """Find reg addrs whose qword content == any vtable slot (object vptr)."""
    hits = []
    mem = os.open(f'/proc/{host}/mem', os.O_RDONLY)
    pats = [struct.pack('<Q', s) for s in slots]
    try:
        for lo, hi in heap_regions(host, excl_start, excl_end):
            try:
                d = os.pread(mem, hi - lo, lo)
            except Exception:
                continue
            if not d:
                continue
            for p in pats:
                i = 0
                while True:
                    i = d.find(p, i)
                    if i < 0:
                        break
                    hits.append(lo + i)
                    i += 1
    finally:
        os.close(mem)
    return hits


def f32(d, off):
    return struct.unpack_from('<f', d, off)[0] if len(d) >= off + 4 else None


def verify(host, base):
    d = mem_read(host, base, 0x541000)
    if d is None or len(d) < 0x540900:
        return None
    f = {
        '0x24': f32(d, 0x24),
        '0x540874': f32(d, 0x540874),
        '0x54087c': f32(d, 0x54087c),
        '0x540880': f32(d, 0x540880),
        '0x540884': f32(d, 0x540884),
        '0x540888': f32(d, 0x540888),
        '0x54088c': f32(d, 0x54088c),
        '0x540894': f32(d, 0x540894),
    }
    ok = (f['0x24'] == 44100.0 and f['0x540874'] == 1.0 and f['0x540894'] == 1.0)
    return f if ok else None


def dump_floats(host, addr, count, tag, od, fmt=lambda v: f'{v:.17g}'):
    d = mem_read(host, addr, count * 4)
    if not d:
        return
    open(f'{tag}.bin', 'wb').write(d)
    n = len(d) // 4
    vals = struct.unpack('<%df' % n, d[:n * 4])
    with open(f'{od}/{tag}.txt', 'w') as f:
        for i, v in enumerate(vals):
            f.write(f'{i}: {fmt(v)}\n')


def main():
    RPP = sys.argv[1] if len(sys.argv) > 1 else '/home/m/soothe-bt/dual300.rpp'
    OD = sys.argv[2] if len(sys.argv) > 2 else '/tmp/rtobj_dual300'
    os.makedirs(OD, exist_ok=True)

    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(f'{OD}/reaper.log', 'w'), stderr=subprocess.STDOUT)
    print(f'reaper pid={proc.pid}', flush=True)

    host = find_host(proc)
    if not host:
        proc.kill(); sys.exit('no host')
    base = soothe2_base(host)
    print(f'host={host} modbase=0x{base:x}', flush=True)

    slots = read_vslots(host, base)
    print(f'vtable slots: {[hex(s) for s in slots]}', flush=True)

    # periodic scan while render alive
    got = None
    t0 = time.time()
    last_host = host
    nscan = 0
    while time.time() - t0 < 300 and not got:
        if proc.poll() is not None:
            print('render finished', flush=True)
            break
        if not os.path.exists(f'/proc/{last_host}'):
            print('host died, re-finding...', flush=True)
            last_host = find_host(proc, wait=10)
            if not last_host:
                break
            base = soothe2_base(last_host)
            if not base:
                break
            slots = read_vslots(last_host, base)
        t1 = time.time()
        hits = scan_objects(last_host, slots, excl_start=base, excl_end=base + 0x7000000)
        if not hits:
            # fallback: no vptr-of-this module found anywhere -> try ctor as marker too
            hits = scan_objects(last_host, [base + 0x529610], excl_start=base, excl_end=base + 0x7000000)
            if hits:
                print(f'   (ctor-marker fallback hits={len(hits)})', flush=True)
        for a in hits:
            f = verify(last_host, a)
            if f:
                got = (a, f)
                break
        nscan += 1
        if not got:
            dt = time.time() - t1
            print(f'  scan#{nscan} t={dt:.2f}s no object', flush=True)
            time.sleep(0.25)
        else:
            break
    host = last_host
    if not got:
        print('NO VALID DSP OBJECT FOUND', flush=True)
        proc.kill(); sys.exit(2)

    addr, fields = got
    tag = os.path.basename(f'obj_{addr:x}')
    meta = {'base': addr, 'modbase': base, 'rpp': RPP, 'host': host, 'fields': fields}
    open(f'{OD}/{tag}.json', 'w').write(json.dumps(meta, indent=2))
    print(f'OBJECT 0x{addr:x} fields={fields}', flush=True)

    dump_floats(host, addr + 0x28, 0x200, f'{tag}.twin_order_f', OD, fmt=lambda v: f'{v:.9g}')
    dump_floats(host, addr + 0x40, 0x2ac, f'{tag}.twin_state_A', OD)
    dump_floats(host, addr + 0x58, 0x2ac, f'{tag}.twin_state_B', OD)
    dump_floats(host, addr + 0x540658, N_DUMP, f'{tag}.window_540658', OD, fmt=lambda v: f'{v:.9g}')
    for off, nm in ((0x5406b8, 'w_b8'), (0x5406c8, 'w_c8'),
                    (0x5406d8, 'w_d8'), (0x5406e8, 'w_e8')):
        dump_floats(host, addr + off, N_DUMP, f'{tag}.{nm}', OD, fmt=lambda v: f'{v:.9g}')
    dump_floats(host, addr + 0x5406a8, 0x200, f'{tag}.warp_5406a8', OD, fmt=lambda v: f'{v:.9g}')
    dump_floats(host, addr + 0x540698, 0x200, f'{tag}.freqaxis_540698', OD, fmt=lambda v: f'{v:.9g}')
    dump_floats(host, addr + 0x5407c8, N_DUMP, f'{tag}.maskacc_5407c8', OD, fmt=lambda v: f'{v:.9g}')

    # scalar fields
    for off, nm in ((0x1a0, 'n1a0'), (0x1ac, 'n1ac'), (0x1a4, 'n1a4'),
                    (0x540868, 'n540868'), (0x54086c, 'n54086c'),
                    (0x540870, 's540870'), (0x540878, 's540878'), (0x54087c, 's54087c'),
                    (0x540880, 's540880'), (0x540884, 's540884'),
                    (0x540888, 's540888'), (0x54088c, 's54088c'),
                    (0x540890, 's540890'), (0x540894, 's540894')):
        d = mem_read(host, addr + off, 4)
        if d:
            v = struct.unpack('<f', d)[0] if not nm.startswith('n') else struct.unpack('<I', d)[0]
            open(f'{OD}/{tag}.{nm}.txt', 'w').write(f'{v!r}\n')

    print('capture done', flush=True)
    proc.wait(timeout=30)


if __name__ == '__main__':
    main()