#!/usr/bin/env python3
"""fntrace2.py — fast single-point INT3 tracer for the live yabridge host.

Flow: poll for host spawn (5ms), immediately PTRACE_SEIZE the main thread,
poke INT3 at FN (no stop required - word write is atomic), then serve
waitpid events (clone children are auto-traced and continued). Logs args at
each hit. Works best against `reaper -renderproject` where all DSP work
happens in a burst right after host spawn.

Usage: python3 scripts/fntrace2.py [rpp] [nhits] [--render]
Env:   FN_ADDR (default 0x180529fe0)
"""
import ctypes
import glob
import json
import os
import signal
import struct
import subprocess
import sys
import time

FN = int(os.environ.get('FN_ADDR', '0x180529FE0'), 16)
SNAPDIR = '/tmp/opencode/fntrace'

PTRACE_CONT = 7
PTRACE_GETREGS = 12
PTRACE_SETREGS = 13
PTRACE_PEEKDATA = 2
PTRACE_POKETEXT = 4
PTRACE_SINGLESTEP = 9
PTRACE_DETACH = 17
PTRACE_SEIZE = 0x4206
PTRACE_INTERRUPT = 0x4207
PTRACE_O_TRACECLONE = 0x00000002


class UserRegs(ctypes.Structure):
    _fields_ = [(n, ctypes.c_ulonglong) for n in (
        'r15', 'r14', 'r13', 'r12', 'rbp', 'rbx', 'r11', 'r10',
        'r9', 'r8', 'rax', 'rcx', 'rdx', 'rsi', 'rdi', 'orig_rax',
        'rip', 'cs', 'eflags', 'rsp', 'ss', 'fs_base', 'gs_base',
        'ds', 'es', 'fs', 'gs')]


libc = ctypes.CDLL('libc.so.6', use_errno=True)


def ptrace(req, pid, addr=0, data=0):
    libc.ptrace.restype = ctypes.c_long
    r = libc.ptrace(req, pid, ctypes.c_void_p(addr), ctypes.c_void_p(data))
    if r == -1:
        e = ctypes.get_errno()
        if req not in (PTRACE_PEEKDATA,):
            raise OSError(e, f'ptrace({req:#x},{pid}) failed')
        return None
    return r


def find_host():
    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'
    nhits = int(sys.argv[2]) if len(sys.argv) > 2 else 12
    render = '--render' in sys.argv
    os.makedirs(SNAPDIR, exist_ok=True)
    subprocess.run('pkill -9 -x reaser 2>/dev/null; pkill -9 -x reaper 2>/dev/null; '
                   "pkill -9 -f '[y]abridge' 2>/dev/null; sleep 1", shell=True)
    if render:
        proc = subprocess.Popen(
            ['/usr/bin/reaper', '-nosplash', '-ignoreerrors',
             '-renderproject', rpp],
            stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT)
    else:
        proc = subprocess.Popen(
            ['/usr/bin/reaper', '-nosplash', '-ignoreerrors', rpp,
             '/home/m/re-tools/play_loop.lua'],
            stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT)

    host = None
    t0 = time.time()
    while time.time() - t0 < 60 and not host:
        host = find_host()
        time.sleep(0.005)
    if not host:
        print('NO HOST')
        return 1
    print('host %d at %.2fs' % (host, time.time() - t0), flush=True)

    ptrace(PTRACE_SEIZE, host, 0, PTRACE_O_TRACECLONE)
    try:
        ptrace(PTRACE_INTERRUPT, host)
        os.waitpid(host, os.WUNTRACED)
    except (OSError, ChildProcessError) as e:
        print('interrupt fail', e)
    orig = ptrace(PTRACE_PEEKDATA, host, FN)
    cc = (orig & ~0xFF) | 0xCC
    ptrace(PTRACE_POKETEXT, host, FN, cc)
    print('armed %#x orig=%#x' % (FN, orig), flush=True)
    try:
        ptrace(PTRACE_CONT, host, 0, 0)
    except OSError:
        pass

    known = {host}
    log = []
    hits = 0
    t_last = time.time()
    deadline_idle = float(os.environ.get('FNTRACE_IDLE', '25'))
    while hits < nhits and time.time() - t_last < deadline_idle:
        try:
            pid, status = os.waitpid(-1, os.WSTOPPED | os.WUNTRACED)
        except ChildProcessError:
            break
        sig = status >> 8
        ev = status >> 16
        known.add(pid)
        t_last = time.time()
        if ev == 3 or ev == 1:      # CLONE/FORK event stop
            try:
                ptrace(PTRACE_CONT, pid, 0, 0)
            except OSError:
                pass
            continue
        if not os.WIFSTOPPED(pid):
            continue
        if sig == signal.SIGTRAP:
            regs = UserRegs()
            try:
                ptrace(PTRACE_GETREGS, pid, 0, ctypes.addressof(regs))
            except OSError:
                continue
            if regs.rip - 1 == FN:
                ret = ptrace(PTRACE_PEEKDATA, pid, regs.rsp)
                rec = dict(ctx=regs.rcx, a2=regs.rdx, cnt=regs.r8 & 0xffffffff,
                           r9=regs.r9 & 0xffffffff, ret=ret, rsp=regs.rsp,
                           rbx=regs.rbx, r12=regs.r12, r13=regs.r13,
                           r14=regs.r14, r15=regs.r15, rsi=regs.rsi, rdi=regs.rdi,
                           rip=regs.rip - 1, tid=pid)
                log.append(rec)
                hits += 1
                print('hit %d tid=%d ctx=%#x a2=%#x cnt=%#x r9d=%#x ret=%#x'
                      % (hits, pid, regs.rcx, regs.rdx,
                         regs.r8 & 0xffffffff, regs.r9 & 0xffffffff, ret),
                      flush=True)
                # step over
                ptrace(PTRACE_POKETEXT, pid, FN, orig)
                regs.rip = FN
                ptrace(PTRACE_SETREGS, pid, 0, ctypes.addressof(regs))
                ptrace(PTRACE_SINGLESTEP, pid, 0, 0)
                try:
                    os.waitpid(pid, os.WUNTRACED)
                except ChildProcessError:
                    pass
                ptrace(PTRACE_POKETEXT, pid, FN, cc)
                try:
                    ptrace(PTRACE_CONT, pid, 0, 0)
                except OSError:
                    pass
            else:
                try:
                    ptrace(PTRACE_CONT, pid, 0, 0)
                except OSError:
                    pass
        else:
            try:
                ptrace(PTRACE_CONT, pid, 0, sig if 0 < sig < 0x20 else 0)
            except OSError:
                pass

    print('hits:', hits, flush=True)
    json.dump(log, open(f'{SNAPDIR}/hits.json', 'w'), indent=1)
    for pid in list(known):
        try:
            os.kill(pid, signal.SIGKILL)
        except OSError:
            pass
    proc.kill()
    print('saved', f'{SNAPDIR}/hits.json')
    return 0


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