#!/usr/bin/env python3
"""
RK3588 DDR Blob Instrumented Tracer
Logs every MMIO read/write with full register context.
Builds a map of: which function accesses which registers in what order.
"""
import struct, sys, json
from unicorn import *
from unicorn.arm64_const import *

MMIO_RANGES = [
    (0xFD580000, 0xFD5A0000, "GRF"),
    (0xFD5F0000, 0xFD600000, "BUS_GRF"),
    (0xFD8C0000, 0xFD8D0000, "SCRU"),
    (0xFE010000, 0xFE030000, "DDRC"),
    (0xFE030000, 0xFE040000, "FW_DDR"),
    (0xFE050000, 0xFE060000, "SGRF"),
    (0xFE0C0000, 0xFE100000, "DDRPHY"),
    (0xFECC0000, 0xFECD0000, "SCRAMBLE"),
    (0xFF000000, 0xFF100000, "SRAM"),
]

def region_name(addr):
    for lo, hi, name in MMIO_RANGES:
        if lo <= addr < hi:
            return f"{name}+0x{addr-lo:X}"
    return f"0x{addr:X}"

def run_trace(blob_path, max_instr=200000, outfile=None):
    with open(blob_path, 'rb') as f:
        blob = f.read()

    trace = []
    count = [0]
    last_pc = [0]
    
    uc = Uc(UC_ARCH_ARM64, UC_MODE_ARM)
    for base in range(0, 0x100000000, 0x10000000):
        try: uc.mem_map(base, 0x10000000, UC_PROT_ALL)
        except: pass
    uc.mem_write(0, blob)
    uc.reg_write(UC_ARM64_REG_SP, 0x00110000 - 16)
    
    # Pre-seed ready values
    for addr, val in [(0xFE0500E0, 0), (0xFE050054, 1)]:
        uc.mem_write(addr, struct.pack('<I', val))
    for ch in [0xFE0C0000, 0xFE0D0000, 0xFE0E0000, 0xFE0F0000]:
        for off, v in [(0xA24, 2), (0x684, 0), (0x10090, 0), (0x10080, 0), (0x10514, 0)]:
            uc.mem_write(ch + off, struct.pack('<I', v))

    def on_mem_read(uc_inst, access, addr, size, value, data):
        is_mmio = any(lo <= addr < hi for lo, hi, _ in MMIO_RANGES)
        if is_mmio:
            val = struct.unpack_from('<I', bytes(uc_inst.mem_read(addr, 4)))[0]
            pc = last_pc[0]
            entry = {
                'op': 'R',
                'addr': f'0x{addr:08X}',
                'reg': region_name(addr),
                'val': f'0x{val:08X}',
                'pc': f'0x{pc:05X}',
                'instr': count[0],
            }
            trace.append(entry)

    def on_mem_write(uc_inst, access, addr, size, value, data):
        is_mmio = any(lo <= addr < hi for lo, hi, _ in MMIO_RANGES)
        if is_mmio:
            pc = last_pc[0]
            entry = {
                'op': 'W',
                'addr': f'0x{addr:08X}',
                'reg': region_name(addr),
                'val': f'0x{value:08X}',
                'pc': f'0x{pc:05X}',
                'instr': count[0],
            }
            trace.append(entry)

    def on_code(uc_inst, addr, size, data):
        count[0] += 1
        last_pc[0] = addr
        if count[0] >= max_instr:
            uc_inst.emu_stop()

    def on_intr(uc_inst, intno, data):
        pc = uc_inst.reg_read(UC_ARM64_REG_PC)
        uc_inst.reg_write(UC_ARM64_REG_PC, pc + 4)

    uc.hook_add(UC_HOOK_MEM_READ, on_mem_read)
    uc.hook_add(UC_HOOK_MEM_WRITE, on_mem_write)
    uc.hook_add(UC_HOOK_CODE, on_code)
    uc.hook_add(UC_HOOK_INTR, on_intr)

    try:
        uc.emu_start(0x10978, len(blob), timeout=60*1000000)
    except:
        pass

    pc = uc.reg_read(UC_ARM64_REG_PC)
    print(f"Executed {count[0]} instructions, final PC=0x{pc:x}")
    print(f"Captured {len(trace)} MMIO accesses")
    
    # Summary
    reads = [t for t in trace if t['op'] == 'R']
    writes = [t for t in trace if t['op'] == 'W']
    print(f"  Reads: {len(reads)}, Writes: {len(writes)}")
    
    # Group by register
    reg_counts = {}
    for t in trace:
        r = t['reg']
        if r not in reg_counts:
            reg_counts[r] = {'R': 0, 'W': 0}
        reg_counts[r][t['op']] += 1
    
    print(f"\nMMIO Register Access Summary:")
    print(f"{'Register':<30s} {'Reads':>6s} {'Writes':>7s}")
    print("-" * 45)
    for reg in sorted(reg_counts.keys()):
        c = reg_counts[reg]
        print(f"{reg:<30s} {c['R']:>6d} {c['W']:>7d}")
    
    # Write full trace
    if outfile:
        with open(outfile, 'w') as f:
            f.write("instr,op,addr,register,value,pc\n")
            for t in trace:
                f.write(f"{t['instr']},{t['op']},{t['addr']},{t['reg']},{t['val']},{t['pc']}\n")
        print(f"\nFull trace written to {outfile}")
    
    # Show first 30 accesses
    print(f"\nFirst 30 MMIO accesses:")
    for t in trace[:30]:
        print(f"  [{t['instr']:6d}] {t['op']} {t['reg']:<25s} = {t['val']}  (PC={t['pc']})")

    return trace

if __name__ == '__main__':
    blob = sys.argv[1] if len(sys.argv) > 1 else '/opt/work/rk3588_ddr_v1.19_trampoline.bin'
    out = sys.argv[2] if len(sys.argv) > 2 else '/opt/work/mmio_trace.csv'
    run_trace(blob, max_instr=200000, outfile=out)