#!/usr/bin/env python3
"""rendersnap.py — SIGSTOP-sampling of the plugin state DURING offline render
(-renderproject): processing is back-to-back, so random stops land inside the
DSP with high probability. Saves every snapshot where FIR != complex identity.
"""
import os
import signal
import struct
import hashlib
import subprocess
import sys
import time

import numpy as np

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
OUT = '/tmp/opencode/rendersnap'


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 300
    os.makedirs(OUT, exist_ok=True)
    subprocess.run("pkill -9 -x reaper; pkill -9 -f '[y]abridge'; sleep 1", shell=True)
    wav = rpp.replace('.rpp', '.wav')
    if os.path.exists(wav):
        os.remove(wav)
    proc = subprocess.Popen(['/usr/bin/reaper', '-nosplash', '-ignoreerrors',
                             '-renderproject', rpp],
                            stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT)
    host = None
    t0 = time.time()
    while time.time() - t0 < 30 and not host:
        host = find_host()
        time.sleep(0.001)
    if not host:
        print('NO HOST')
        return 1
    print('host %d at %.3fs' % (host, time.time() - t0), flush=True)

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

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

    ctx = None
    vt = struct.pack('<Q', 0x1824AC210)
    m48 = struct.pack('<I', 0x473b8000)

    def scan_ctx():
        for line in open(f'/proc/{host}/maps'):
            parts = line.split()
            if 'rw' not in parts[1]:
                continue
            lo, hi = (int(x, 16) for x in parts[0].split('-'))
            CH = 16 * 1024 * 1024
            a = lo
            while a < hi:
                d = rd(a, min(CH + 4096, hi - a))
                if not d:
                    break
                j = d.find(vt)
                while j >= 0:
                    cand = a + j
                    sb = rd(cand + 0x540870, 4)
                    if sb and struct.unpack('<f', sb)[0] > 100:
                        return cand
                    j = d.find(vt, j + 1)
                j = d.find(m48)
                while j >= 0:
                    cand = a + j - 0x24
                    sb = rd(cand + 0x540870, 4)
                    if sb and struct.unpack('<f', sb)[0] > 100:
                        return cand
                    j = d.find(m48, j + 1)
                a += CH
        return None

    # find-loop: stop-scan-resume until instance exists (render lasts ~0.5s)
    while ctx is None and time.time() - t0 < 25:
        try:
            os.kill(host, signal.SIGSTOP)
        except ProcessLookupError:
            break
        try:
            ctx = scan_ctx()
        finally:
            if ctx is None:
                try:
                    os.kill(host, signal.SIGCONT)
                except ProcessLookupError:
                    break
        if ctx is None:
            time.sleep(0.004)
    print('ctx %#x' % ctx if ctx else 'NO CTX', flush=True)
    if not ctx:
        return 1
    print('ctx %#x' % ctx if ctx else 'NO CTX', flush=True)
    if not ctx:
        return 1

    hits = saved = 0
    rng = np.random.default_rng(3)
    prev_sig = None
    phases = []
    while True:
        try:
            os.kill(host, signal.SIGSTOP)
        except ProcessLookupError:
            break
        try:
            pb = rd(ctx + 0x540668, 8)
            if not pb:
                continue
            p = struct.unpack('<Q', pb)[0]
            fb = rd(p, NARR * 4)
            if not fb:
                continue
            arr = np.frombuffer(fb[:2049 * 8], dtype='<f4').astype(np.float32)
            sig = arr.tobytes()[:4096]
            rb = rd(struct.unpack('<Q', rd(ctx + 0x5407f8, 8))[0], 2049 * 4)
            rsig = rb[:512] if rb else b''
            key = hashlib.md5(sig + rsig).digest()
            if key != prev_sig:
                prev_sig = key
                mag = np.hypot(arr[0::2], arr[1::2])
                rv = np.frombuffer(rb, dtype='<f4') if rb else None
                phase = dict(t=round(time.time() - t0, 3),
                             firMax=float(np.abs(arr[1:2049]).max()),
                             fir43=float(mag[43]), fir171=float(mag[171]),
                             r43=float(rv[43]) if rv is not None else -1,
                             r171=float(rv[171]) if rv is not None else -1)
                phases.append(phase)
                store = {}
                for off in SLOTS:
                    q = rd(ctx + off, 8)
                    if not q:
                        continue
                    ptr = struct.unpack('<Q', q)[0]
                    if ptr < 0x10000:
                        continue
                    ab = rd(ptr, NARR * 4)
                    if ab:
                        store[hex(off)] = np.frombuffer(ab, dtype='<f4').astype(np.float32)
                fn = f'{OUT}/phase{saved:03d}.npz'
                np.savez_compressed(fn, **store)
                saved += 1
                print('PHASE %s -> %s' % (phase, fn), flush=True)
                if saved >= 24:
                    break
        finally:
            try:
                os.kill(host, signal.SIGCONT)
            except ProcessLookupError:
                pass
        time.sleep(float(rng.uniform(0.001, 0.01)))
        try:
            os.kill(host, 0)
        except ProcessLookupError:
            print('host exited at %.2fs' % (time.time() - t0), flush=True)
            break
    print('phases=%d' % saved)
    os.close(fd)
    proc.kill()
    return 0


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