#!/usr/bin/env python3
"""wine_chain_trace.py — живой захват промежуточных состояний FIR-цепи
soothe2 через winedbg (wine) + /proc/<pid>/mem.

Брейкпоинты:
  EXP  0x1803831c0  комплексная экспонента FIR-цепи (rcx=buf, r8d=count float)
  DF0  0x18000b3c0  финальный complex-mul (rcx=FIR, rdx=track, r8d=n пар)
На хите: читаем rcx/rdx/r8 (info reg), буферы — через /proc/<pid>/mem,
копим сэмплы, отпускаем (c). Рендер не убивается.

Запуск: python3 scripts/wine_chain_trace.py <rpp> [n_hits] [outdir]
"""
import os
import pickle
import re
import signal
import struct
import subprocess
import sys
import threading
import time

import numpy as np

BP_EXP = 0x1803831c0
BP_DF0 = 0x18000b3c0
CTX_SLOTS = {'scr': 0x540628, 'trk': 0x540688, 'cur': 0x540678,
             'fir_ptr': 0x540668}


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, cmd[:80]
    return None, None


def find_ctx(fd, pid):
    vt = struct.pack('<Q', 0x1824AC210)
    m48 = struct.pack('<I', 0x47380000)
    for line in open(f'/proc/{pid}/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:
            n = min(CH, hi - a)
            try:
                d = os.pread(fd, n, a)
            except OSError:
                break
            j = d.find(vt)
            while j >= 0:
                cand = a + j
                sb = os.pread(fd, 4, cand + 0x540870)
                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
                try:
                    sb = os.pread(fd, 4, cand + 0x540870)
                    if sb and struct.unpack('<f', sb)[0] > 100:
                        return cand
                except OSError:
                    pass
                j = d.find(m48, j + 1)
            a += n
    return None


class WineDbg:
    """Асинхронный ридер stdout winedbg + обмен командами по приглашению."""

    PROMPT = 'Wine-dbg>'

    def __init__(self, pid):
        self.p = subprocess.Popen(
            ['winedbg', '--pid', str(pid)],
            stdin=subprocess.PIPE, stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT, text=True, bufsize=1)
        self.buf = ''
        self.lock = threading.Lock()
        self.ev = threading.Event()
        self.alive = True
        self.t = threading.Thread(target=self._reader, daemon=True)
        self.t.start()
        if not self.ev.wait(30):
            raise TimeoutError('winedbg не показал приглашение')

    def _reader(self):
        while self.alive:
            ch = self.p.stdout.read(1)
            if not ch:
                self.alive = False
                self.ev.set()
                return
            with self.lock:
                self.buf += ch
                if self.PROMPT in self.buf:
                    self.ev.set()

    def cmd(self, c, timeout=90):
        with self.lock:
            self.buf = ''
            self.ev.clear()
        self.p.stdin.write(c + '\n')
        self.p.stdin.flush()
        if not self.ev.wait(timeout):
            with self.lock:
                tail = self.buf[-300:]
            raise TimeoutError('winedbg timeout после %r; tail=%r' % (c, tail))
        with self.lock:
            out = self.buf.replace(self.PROMPT, '').strip()
            self.buf = ''
            self.ev.clear()
        return out

    def close(self):
        self.alive = False
        try:
            self.p.stdin.write('quit\n')
            self.p.stdin.flush()
        except Exception:
            pass
        try:
            self.p.kill()
        except Exception:
            pass


def parse_regs(text):
    regs = {}
    for mm in re.finditer(r'\b([re]?[a-z]{2,3}|r\d+d?)\s*[:=]\s*([0-9a-fA-F]{4,16})\b', text):
        name = mm.group(1).lower()
        val = int(mm.group(2), 16)
        if name not in regs:
            regs[name] = val
    # нормализация имён к 64-битным
    alias = {'eax': 'rax', 'ecx': 'rcx', 'edx': 'rdx', 'ebx': 'rbx',
             'esi': 'rsi', 'edi': 'rdi', 'ebp': 'rbp', 'esp': 'rsp'}
    out = {}
    for k, v in regs.items():
        k64 = alias.get(k, k)
        if k64.startswith('r') and k64.endswith('d') and k64[1:-1].isdigit():
            k64 = k64[:-1]
        if len(k64) <= 3 or k64.startswith('r'):
            out[k64] = v
    return out


def main():
    rpp = sys.argv[1] if len(sys.argv) > 1 else '/home/m/soothe-bt/dual_b1q_0.5.rpp'
    n_target = int(sys.argv[2]) if len(sys.argv) > 2 else 60
    outdir = sys.argv[3] if len(sys.argv) > 3 else '/tmp/opencode/winetrace'
    os.makedirs(outdir, exist_ok=True)

    wav = None
    for ln in open(rpp, 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)

    subprocess.run("pkill -9 -x reaper; pkill -9 -f '[y]abridge'; "
                   "rm -rf /run/user/1000/yabridge-soothe2_x64-*; sleep 1",
                   shell=True)
    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, cmdl = find_host()
        if not host:
            time.sleep(0.002)
    if not host:
        print('NO HOST')
        return 1
    print('host %d (%s)' % (host, cmdl), flush=True)
    fd = os.open(f'/proc/{host}/mem', os.O_RDONLY)

    ctx = None
    while ctx is None and time.time() - t0 < 25:
        try:
            os.kill(host, signal.SIGSTOP)
        except ProcessLookupError:
            break
        ctx = find_ctx(fd, host)
        os.kill(host, signal.SIGCONT)
        if not ctx:
            time.sleep(0.005)
    if not ctx:
        print('NO CTX')
        return 1
    print('ctx %#x' % ctx, flush=True)

    dbg = WineDbg(host)
    print(dbg.cmd('break *%#x' % BP_EXP)[:160], flush=True)
    print(dbg.cmd('break *%#x' % BP_DF0)[:160], flush=True)

    def rd(a, n):
        return os.pread(fd, n, a)

    def rd_f32(a, n):
        return np.frombuffer(rd(a, 4*n), dtype='<f4').astype(np.float64)

    def rd_q(a):
        return struct.unpack('<Q', rd(a, 8))[0]

    samples = []
    hits = {'EXP': 0, 'DF0': 0}
    t_start = time.time()
    stall = 0
    while sum(hits.values()) < n_target and time.time() - t_start < 300:
        try:
            out = dbg.cmd('c', timeout=120)
        except TimeoutError as e:
            print('timeout:', str(e)[-200:], flush=True)
            stall += 1
            if stall >= 3:
                break
            continue
        addrs = [int(x, 16) for x in re.findall(r'0x[0-9a-fA-F]{9,}', out)]
        pc = None
        for a in addrs:
            if abs(a - BP_EXP) < 64:
                pc = a; kind = 'EXP'; break
            if abs(a - BP_DF0) < 64:
                pc = a; kind = 'DF0'; break
        if pc is None:
            ir = dbg.cmd('info reg', timeout=30)
            rr = parse_regs(ir)
            pc = rr.get('rip', 0)
            kind = 'EXP' if abs(pc-BP_EXP) < 64 else ('DF0' if abs(pc-BP_DF0) < 64 else None)
            if kind is None:
                stall += 1
                if stall >= 5:
                    print('неопознанные остановки; tail:', out[-200:], flush=True)
                    break
                continue
        ir = dbg.cmd('info reg', timeout=30)
        rr = parse_regs(ir)
        rcx = rr.get('rcx', 0); rdx = rr.get('rdx', 0); r8 = rr.get('r8', 0)
        rec = {'kind': kind, 'rip': pc, 'rcx': rcx, 'rdx': rdx, 'r8': r8,
               't': round(time.time()-t_start, 4)}
        try:
            if kind == 'EXP':
                rec['buf'] = rd_f32(rcx, 4098)
                rec['count'] = r8
            else:
                rec['fir'] = rd_f32(rcx, 4098)
                if rdx > 0x10000:
                    rec['track'] = rd_f32(rdx, 2049*2)
            # слоты контекста тем же мгновением (процесс остановлен!)
            rec['scr'] = rd_f32(ctx+CTX_SLOTS['scr'], 2049)
            rec['trk'] = rd_f32(ctx+CTX_SLOTS['trk'], 2049)
            rec['cur'] = rd_f32(ctx+CTX_SLOTS['cur'], 2049)
            fp = rd_q(ctx+CTX_SLOTS['fir_ptr'])
            rec['fir_via_ctx'] = rd_f32(fp, 4098)
        except OSError as e:
            rec['err'] = str(e)
        samples.append(rec)
        hits[kind] += 1
        if sum(hits.values()) % 10 == 0:
            print('hits:', hits, flush=True)

    print('сбор завершён:', hits, flush=True)
    snap_ptrs = {}
    snap_arr = {}
    for nm, off in CTX_SLOTS.items():
        try:
            p = rd_q(ctx+off)
            if p > 0x10000:
                snap_ptrs[nm] = p
                snap_arr[nm] = rd_f32(p, 4100)
        except OSError:
            pass
    dbg.close()

    with open(os.path.join(outdir, 'chain_samples.pkl'), 'wb') as f:
        pickle.dump({'samples': samples, 'snap_ptrs': snap_ptrs, 'ctx': ctx}, f)
    np.savez_compressed(os.path.join(outdir, 'ctx_snap.npz'), **snap_arr)
    print('saved', len(samples), 'samples ->', outdir, flush=True)
    for _ in range(600):
        if proc.poll() is not None:
            break
        time.sleep(0.1)
    print('reaper_rc=%s wav=%s' % (proc.poll(),
          os.path.getsize(wav) if wav and os.path.exists(wav) else 'NONE'), flush=True)
    return 0


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