#!/usr/bin/env python3
"""scan3.py — Multi-instance DSP context scanner for soothe2.

Fixes over scan2.py:
1. Pre-scans ALL processes for ctx (no host-finding delay)
2. Scans for ALL instances (GUI/DSP pair hypothesis from 24c)
3. Captures full state per instance for comparison
4. Uses rendersnap-style sampling for FIR/slot captures
"""
import glob
import hashlib
import os
import signal
import struct
import subprocess
import sys
import time

import numpy as np

OUT = '/tmp/opencode/scan3'
NARR = 8194

SLOTS_FULL = [
    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,
]

VTQ = struct.pack('<Q', 0x1824AC210)
M48 = struct.pack('<I', 0x473b8000)


def pre_scan_all(max_region=50*1024*1024):
    """Scan ALL processes for DSP ctx instances (no host needed)."""
    instances = {}  # pid -> [(ctx_addr, sens)]
    for p in glob.glob('/proc/[0-9]*'):
        pid = int(os.path.basename(p))
        try:
            maps = open(f'/proc/{pid}/maps').read()
        except Exception:
            continue
        if 'soothe2' not in maps:
            continue
        try:
            fd = os.open(f'/proc/{pid}/mem', os.O_RDONLY)
        except Exception:
            continue
        pid_insts = []
        for line in maps.split('\n'):
            parts = line.split()
            if len(parts) < 2 or 'rw' not in parts[1]:
                continue
            lo, hi = (int(x, 16) for x in parts[0].split('-'))
            if hi - lo > max_region:
                continue
            try:
                data = os.pread(fd, min(hi - lo, 4*1024*1024), lo)
            except Exception:
                continue
            for pat, off in ((VTQ, 0), (M48, -0x24)):
                j = data.find(pat)
                while j >= 0:
                    cand = lo + j + off
                    try:
                        sb = os.pread(fd, 4, cand + 0x540870)
                        if sb and struct.unpack('<f', sb)[0] > 100:
                            pid_insts.append(cand)
                    except Exception:
                        pass
                    j = data.find(pat, j + 1)
        os.close(fd)
        if pid_insts:
            instances[pid] = list(set(pid_insts))
    return instances


def read_state(fd, ctx):
    """Read key DSP state from a context."""
    state = {'ctx': ctx}
    for off, name, fmt in [
        (0x540870, 'sens', '<f'),
        (0x540874, 'depth', '<f'),
        (0x54087c, 'mix', '<f'),
        (0x540888, 'att_coeff', '<f'),
        (0x54088c, 'rel_coeff', '<f'),
        (0x1a0, 'nfft', '<i'),
    ]:
        try:
            sb = os.pread(fd, 4, ctx + off)
            state[name] = struct.unpack(fmt, sb)[0]
        except Exception:
            state[name] = None

    try:
        pb = os.pread(fd, 8, ctx + 0x540668)
        p = struct.unpack('<Q', pb)[0]
        if p > 0x10000:
            fb = os.pread(fd, 2049*8, p)
            arr = np.frombuffer(fb, dtype='<f4')
            mag = np.hypot(arr[0::2], arr[1::2])
            state['fir_mag0'] = float(mag[0])
            state['fir_mag43'] = float(mag[43]) if len(mag) > 43 else -1
            state['fir_mag171'] = float(mag[171]) if len(mag) > 171 else -1
            state['fir_is_identity'] = bool(np.all(np.abs(mag[:50] - 1.0) < 0.01))
    except Exception:
        pass
    return state


def sampling_phase(fd, host, ctx, t_start):
    """Rendersnap-style FIR sampling."""
    rng = np.random.default_rng(3)
    prev_sig = None
    saved = 0
    while time.time() - t_start < 20:
        try:
            os.kill(host, signal.SIGSTOP)
        except ProcessLookupError:
            break
        try:
            pb = os.pread(fd, 8, ctx + 0x540668)
            p = struct.unpack('<Q', pb)[0]
            if p < 0x10000:
                continue
            fb = os.pread(fd, NARR * 4, p)
            arr = np.frombuffer(fb[:2049*8], dtype='<f4').astype(np.float32)
            sig = arr.tobytes()[:4096]

            rb_ptr = os.pread(fd, 8, ctx + 0x5407f8)
            rp = struct.unpack('<Q', rb_ptr)[0]
            rb = os.pread(fd, 2049*4, rp) if rp > 0x10000 else None
            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() - t_start, 3),
                    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,
                )
                store = {}
                for off in SLOTS_FULL:
                    try:
                        q = os.pread(fd, 8, ctx + off)
                        ptr = struct.unpack('<Q', q)[0]
                        if ptr > 0x10000:
                            ab = os.pread(fd, NARR * 4, ptr)
                            store[hex(off)] = np.frombuffer(ab, dtype='<f4').astype(np.float32)
                    except Exception:
                        pass
                fn = f'{OUT}/scan{saved:03d}.npz'
                np.savez_compressed(fn, **store)
                saved += 1
                print(f'  PHASE {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(f'  host exited at {time.time()-t_start:.2f}s')
            break
    return saved


def main():
    rpp = sys.argv[1] if len(sys.argv) > 1 else '/home/m/soothe-bt/dual_b1q_0.5.rpp'
    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)

    t0 = time.time()
    all_instances = {}  # pid -> [ctx_addrs]

    # Phase 1: aggressive pre-scan (no host needed)
    print('Pre-scanning for DSP contexts...', flush=True)
    for att in range(100):
        found = pre_scan_all()
        for pid, ctxs in found.items():
            if pid not in all_instances:
                all_instances[pid] = ctxs
                print(f'  pid={pid} ctx={[hex(c) for c in ctxs]} at {time.time()-t0:.3f}s', flush=True)
        if all_instances:
            break
        time.sleep(0.0002)

    if not all_instances:
        print('NO CTX FOUND')
        proc.kill()
        return 1

    # Phase 2: read state of each instance
    for pid, ctxs in all_instances.items():
        for ctx in ctxs:
            fd = os.open(f'/proc/{pid}/mem', os.O_RDONLY)
            state = read_state(fd, ctx)
            os.close(fd)
            print(f'\n=== Instance pid={pid} ctx={hex(ctx)} ===')
            for k, v in sorted(state.items()):
                if isinstance(v, float):
                    print(f'  {k}: {v:.6f}')
                else:
                    print(f'  {k}: {v}')

    # Phase 3: rendersnap-style sampling on the primary
    primary_pid = min(all_instances.keys())
    primary_ctx = min(all_instances[primary_pid])
    print(f'\n--- Sampling primary {hex(primary_ctx)} (pid={primary_pid}) ---')

    fd = os.open(f'/proc/{primary_pid}/mem', os.O_RDONLY)
    saved = sampling_phase(fd, primary_pid, primary_ctx, time.time())
    os.close(fd)

    print(f'\ntotal samples={saved}')
    proc.kill()
    return 0


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