#!/usr/bin/env python3
"""rendersnap2.py — non-destructive variant of rendersnap.py.

Differences (24j):
- does NOT kill the host: lets -renderproject finish naturally so the output
  wav is complete (partial-wav hazard, NOTES 24c);
- additionally dumps the ctx SCALAR bank 0x540860..0x5408c0 every phase
  (goal: wet/FIR scalar [ctx+0x540888], BLOCKMAP 23b line 52b832);
- keeps sampling until host exit or phase cap, dedupe by FIR+scalar md5.
"""
import hashlib
import os
import signal
import struct
import subprocess
import sys
import time

import numpy as np

SLOTS = [0x540668, 0x540548, 0x540550, 0x540598, 0x540628, 0x5406f8,
         0x540678, 0x540688, 0x5406c8, 0x5406e8, 0x540768,
         0x540788, 0x5407f8,
         # 24mm9: ACC-таблица указателей (шаг 10 combine) и WINfreq
         # (окно FIR-цепи; падающий Hann — контроль формы)
         0x5407c8, 0x540658]
NARR = 8194
SCAL_OFF = 0x540860
SCAL_N = 24  # floats -> 0x540860..0x5408c0
OUT = '/tmp/opencode/rendersnap2'


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'
    cap = int(sys.argv[2]) if len(sys.argv) > 2 else 200
    global OUT
    OUT = sys.argv[3] if len(sys.argv) > 3 else OUT
    os.makedirs(OUT, exist_ok=True)
    subprocess.run("pkill -9 -x reaper; pkill -9 -f '[y]abridge'; "
                   "rm -rf /run/user/1000/yabridge-soothe2_x64-*; sleep 1", shell=True)
    # wav path = RENDER_FILE from the project (NOT rpp.replace!) — deleting the
    # wrong file destroyed corpus refs once (24k-3 hazard).
    wav = None
    for ln in open(rpp, 'r', errors='replace'):
        if 'RENDER_FILE' in ln and '"' in ln:
            wav = ln.split('"')[1]
            break
    if wav and 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)
    t0 = time.time()
    host = None
    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

    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

    ctx = None
    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)
    if not ctx:
        print('NO CTX')
        return 1
    print('ctx %#x' % ctx, flush=True)

    rng = np.random.default_rng(11)
    prev = None
    saved = 0
    while True:
        try:
            os.kill(host, signal.SIGSTOP)
        except ProcessLookupError:
            print('host exited at %.2fs' % (time.time() - t0), flush=True)
            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
            scal = rd(ctx + SCAL_OFF, SCAL_N * 4)
            if scal is None:
                continue
            key = hashlib.md5(fb[:2049 * 8] + scal).digest()
            if key != prev:
                prev = key
                arr = np.frombuffer(fb[:2049 * 8], dtype='<f4')
                mag = np.hypot(arr[0::2], arr[1::2])
                sv = np.frombuffer(scal, dtype='<f4')
                store = {'fir_re': arr[0::2].copy(), 'fir_im': arr[1::2].copy(),
                         'scal': sv.copy()}
                # LUT params: sub-object pointer at [ctx+0x180] (24z: 563a60 uses
                # r15=rcx -> [r15+0x180]); A=[sub+0] B=[sub+4] G=[sub+0xc]
                sb = rd(ctx + 0x180, 8)
                if sb:
                    sub = struct.unpack('<Q', sb)[0]
                    if sub > 0x10000:
                        fb4 = rd(sub, 0x18)
                        if fb4:
                            store['lut_A'] = struct.unpack('<f', fb4[0:4])[0]
                            store['lut_B'] = struct.unpack('<f', fb4[4:8])[0]
                            store['lut_G'] = struct.unpack('<f', fb4[12:16])[0]
                            store['lut_mode'] = fb4[16]
                for off in SLOTS[1:]:
                    if off == 0x5407c8:
                        # ACC is table of pointers per-band [5407c8 + i*8] -> buf 2049 f32 each
                        for bi in range(4):
                            q = rd(ctx + off + 8*bi, 8)
                            if not q: continue
                            ptr = struct.unpack('<Q', q)[0]
                            if ptr < 0x10000: continue
                            ab = rd(ptr, 2049 * 4)
                            if ab:
                                store[hex(off) + f"_{bi}"] = np.frombuffer(ab, dtype='<f4').astype(np.float32)
                        # also store raw table pointer for debug
                        q = rd(ctx + off, 8)
                        if q:
                            ptr = struct.unpack('<Q', q)[0]
                            if ptr >= 0x10000:
                                ab = rd(ptr, NARR * 4)
                                if ab:
                                    store[hex(off)] = np.frombuffer(ab, dtype='<f4').astype(np.float32)
                        continue
                    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)
                np.savez_compressed(f'{OUT}/ph{saved:03d}.npz', t_snap=time.time() - t0,
                                    **store)
                print('PH%03d t=%.2f fir43=%.4f fir171=%.4f | s888=%.6f s88c=%.6f s874=%.6f s87c=%.6f s870=%.3f' %
                      (saved, time.time() - t0, mag[43], mag[171],
                       sv[10], sv[11], sv[5], sv[7], sv[4]), flush=True)
                saved += 1
                if saved >= cap:
                    break
        finally:
            try:
                os.kill(host, signal.SIGCONT)
            except ProcessLookupError:
                pass
        time.sleep(float(rng.uniform(0.0005, 0.006)))
        try:
            os.kill(host, 0)
        except ProcessLookupError:
            print('host exited at %.2fs' % (time.time() - t0), flush=True)
            break
    # wait for natural finish so wav completes
    for _ in range(300):
        if proc.poll() is not None:
            break
        time.sleep(0.1)
    print('saved=%d reaper_rc=%s wav=%s' % (saved, proc.poll(),
          os.path.getsize(wav) if wav and os.path.exists(wav) else 'NONE'))
    return 0


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