#!/usr/bin/env python3
"""tp_slot_probe.py — snapshot tp[0x4f] and tp[0x55] at the exact PC
where fn_5540 reads them, and dump tp[base..base+0x2ac] for diff.

Runs on vendor.bin and rebuilt.bin side-by-side. The PCs for the two
ldrs differ between vendor and rebuilt codegen, so we probe multiple
candidate PCs by looking for `ldr w, [x19, #0x154]` and
`ldr w, [x19, #0x13c]` equivalents.
"""
import argparse
import os
import sys
from unicorn import *
from unicorn.arm64_const import *

sys.path.insert(0, os.path.join(
    os.path.dirname(os.path.abspath(__file__)), '..'))
from mmio_diff import (SRAM_BASE, BLOB_BASE, STACK_BASE, RET_STUB,
                       MMIO, XREG, stub_value, reset_stub_state)

# Search window: PCs between first fn_5540 MMIO write (match 194 = 0x6978)
# and the diverging write (0x69c8/0x6a74). Capture x19 + read the two
# slots as soon as we see an instruction pattern `ldr w?, [x19, #0x154]`
# or `ldr w?, [x19, #0x13c]` (encoding 0xb94_154.. / 0xb94_13c..).
# Simpler: hook every PC in fn_5540 [0x5540..0x6040) and, on each, if
# the insn looks like such an ldr, snapshot x19 and memory contents.


def match_ldr_w_imm(ins, imm):
    """Match `ldr w?, [x?, #imm]` — any base reg. Returns rn or None."""
    if (ins >> 22) != 0b1011100101:
        return None
    imm12 = (ins >> 10) & 0xFFF
    if imm12 * 4 != imm:
        return None
    return (ins >> 5) & 0x1F


def run(blob_path, max_insn=500_000):
    reset_stub_state()
    blob = open(blob_path, 'rb').read()
    uc = Uc(UC_ARCH_ARM64, UC_MODE_ARM)
    uc.mem_map(SRAM_BASE, 0x100000, UC_PROT_ALL)
    uc.mem_write(BLOB_BASE, blob)
    uc.mem_map(STACK_BASE, 0x100000, UC_PROT_ALL)
    uc.mem_map(RET_STUB, 0x1000, UC_PROT_ALL)
    uc.mem_write(RET_STUB, b'\x00\x00\x20\xd4')
    for b, s in MMIO:
        uc.mem_map(b, s, UC_PROT_ALL)

    state = {'count': 0, 'prev_pc': 0, 'same_pc': 0,
             'snap': None, 'tp_dump': None}

    def hook_code(uc, addr, size, ud):
        state['count'] += 1
        if addr == state.get('prev_pc'):
            state['same_pc'] += 1
            if state['same_pc'] > 10000: uc.emu_stop()
        else:
            state['same_pc'] = 0; state['prev_pc'] = addr
        if state['count'] >= max_insn: uc.emu_stop()
        # Whole-blob search — the rebuilt's ldr site may be anywhere in
        # fn_5540 post-recompile.
        if not (0xff001000 <= addr < 0xff020000):
            return
        try:
            ins = int.from_bytes(uc.mem_read(addr, 4), 'little')
        except UcError:
            return
        # Look for ldr w, [x?, #0x13c] — this is the tp[0x4f] load
        rn = match_ldr_w_imm(ins, 0x13c)
        if rn is not None and state['snap'] is None:
            base = uc.reg_read(XREG[rn])
            try:
                tp4f = int.from_bytes(uc.mem_read(base + 0x13c, 4), 'little')
                tp55 = int.from_bytes(uc.mem_read(base + 0x154, 4), 'little')
                dump = uc.mem_read(base, 0x2ac)
            except UcError:
                return
            state['snap'] = (addr, base, rn, tp4f, tp55)
            state['tp_dump'] = bytes(dump)

    def hook_mmio_read(uc, typ, addr, size, val, ud):
        v = stub_value(addr) & ((1 << size*8) - 1)
        uc.mem_write(addr, v.to_bytes(size, 'little'))

    def hook_mmio_write(uc, typ, addr, size, val, ud):
        pass

    def hook_unmapped(uc, typ, addr, size, val, ud):
        page = addr & ~0xFFFF
        try: uc.mem_map(page, 0x10000, UC_PROT_ALL)
        except UcError: pass
        if typ == UC_MEM_READ_UNMAPPED:
            v = stub_value(addr) & ((1 << size*8) - 1)
            uc.mem_write(addr, v.to_bytes(size, 'little'))
        return True

    uc.hook_add(UC_HOOK_CODE, hook_code)
    for b, s in MMIO:
        uc.hook_add(UC_HOOK_MEM_READ, hook_mmio_read, begin=b, end=b + s)
        uc.hook_add(UC_HOOK_MEM_WRITE, hook_mmio_write, begin=b, end=b + s)
    uc.hook_add(UC_HOOK_MEM_UNMAPPED, hook_unmapped)

    uc.reg_write(UC_ARM64_REG_SP, STACK_BASE + 0xF0000)
    uc.reg_write(UC_ARM64_REG_X30, BLOB_BASE + 0x40)

    pc = BLOB_BASE
    remaining = max_insn
    while remaining > 0:
        try:
            uc.emu_start(pc, RET_STUB, count=remaining); break
        except UcError as e:
            pc = uc.reg_read(UC_ARM64_REG_PC)
            try: insn = int.from_bytes(uc.mem_read(pc, 4), 'little')
            except UcError: break
            if (insn >> 20) == 0xD53:
                rt = insn & 0x1F
                if rt < 31: uc.reg_write(XREG[rt], 0)
                pc += 4; uc.reg_write(UC_ARM64_REG_PC, pc); remaining -= 1; continue
            if (insn >> 20) in (0xD51, 0xD50):
                pc += 4; uc.reg_write(UC_ARM64_REG_PC, pc); remaining -= 1; continue
            break

    return state


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('blob')
    args = ap.parse_args()
    state = run(args.blob)
    if state['snap'] is None:
        print(f'NO snapshot captured — did not find ldr w, [x19, #0x13c] in fn_5540 range')
        return
    pc, tp_base, rn, tp4f, tp55 = state['snap']
    print(f'fn_5540 read site: pc=0x{pc:x} (via x{rn})')
    print(f'  tp_base=0x{tp_base:x}')
    print(f'  tp[0x4f] (+0x13c) = 0x{tp4f:08x}')
    print(f'  tp[0x55] (+0x154) = 0x{tp55:08x}')
    print(f'  computed write val = 0x{(tp55 | (tp4f << 16)) & 0xFFFFFFFF:08x}')
    import sys as _sys
    print('  tp dump (32-byte-line, u32):')
    dump = state['tp_dump']
    for off in range(0, min(0x2ac, len(dump)), 32):
        import struct
        words = struct.unpack_from('<8I', dump, off)
        print(f'    +0x{off:03x}: ' + ' '.join(f'{w:08x}' for w in words))


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