#!/usr/bin/env python3
"""tp_slot_writes.py — list every write to the tp buffer's +0x13c slot
(tp[0x4f]) during the run, on both vendor and rebuilt. Since tp lives
in SRAM at 0xff0164f8 (discovered via tp_slot_probe), we just hook
SRAM writes to that exact address.
"""
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)

TP_BASE = 0xff0164f8
TP_SLOT_4F = TP_BASE + 0x13c
TP_SLOT_55 = TP_BASE + 0x154


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, 'writes_4f': [],
             'writes_55': []}

    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()

    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_sram_write(uc, typ, addr, size, val, ud):
        if addr == TP_SLOT_4F or (addr <= TP_SLOT_4F < addr + size):
            pc = uc.reg_read(UC_ARM64_REG_PC)
            state['writes_4f'].append((state['count'], pc, addr, size, val))
        if addr == TP_SLOT_55 or (addr <= TP_SLOT_55 < addr + size):
            pc = uc.reg_read(UC_ARM64_REG_PC)
            state['writes_55'].append((state['count'], pc, addr, size, val))

    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)
    # Hook SRAM writes in the whole blob+data region where tp lives.
    uc.hook_add(UC_HOOK_MEM_WRITE, hook_sram_write,
                begin=SRAM_BASE, end=SRAM_BASE + 0x100000)
    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)
    print(f'=== writes to tp[0x4f] (@{TP_SLOT_4F:#x}) — {len(state["writes_4f"])} total ===')
    for tick, pc, addr, size, val in state['writes_4f']:
        off = pc - 0xff001000
        print(f'  tick={tick:7d}  pc=0x{pc:x} (blob+0x{off:05x})  addr=0x{addr:x} sz={size}  val=0x{val:x}')
    print(f'=== writes to tp[0x55] (@{TP_SLOT_55:#x}) — {len(state["writes_55"])} total ===')
    for tick, pc, addr, size, val in state['writes_55']:
        off = pc - 0xff001000
        print(f'  tick={tick:7d}  pc=0x{pc:x} (blob+0x{off:05x})  addr=0x{addr:x} sz={size}  val=0x{val:x}')


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