#!/usr/bin/env python3
"""
Deep MMIO tracer - injects "ready" responses to polls to push execution further.
When a register is read multiple times from the same PC (a poll), inject the
expected value to make the condition pass.
"""
import struct, sys
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 rname(addr):
    for lo, hi, n in MMIO_RANGES:
        if lo <= addr < hi:
            return f"{n}+0x{addr-lo:X}"
    return f"0x{addr:X}"

# Smart responses: after N reads from same PC, inject "ready" value
# These simulate the PHY completing operations
READY_RESPONSES = {
    # SGRF
    0xFE0500E0: 0x00000000,  # status = ready
    0xFE050054: 0x00000001,  # CON21 = done
    # DDRPHY - all channels
}

# PHY register ready values (applied to all 4 channels)
PHY_READY = {
    0x0A24: 0x00000002,  # DfiStatus = ready (bit 1 set)
    0x0684: 0x00000000,  # CalBusy = not busy
    0x10090: 0x00000000, # MicroContMuxSel = available
    0x10080: 0x00000000, # MicroReset = done (not negative)
    0x10514: 0x00000001, # UctWriteProtShadow = done (bit 0 set)
    0x10510: 0x00000000, # UctWriteOnlyShadow
}

# Build full ready map
for ch_base in [0xFE0C0000, 0xFE0D0000, 0xFE0E0000, 0xFE0F0000]:
    for off, val in PHY_READY.items():
        READY_RESPONSES[ch_base + off] = val

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

    trace = []
    count = [0]
    last_pc = [0]
    poll_counts = {}  # (pc, addr) -> count
    injected = [0]
    unique_pcs = set()
    
    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 all ready values
    for addr, val in READY_RESPONSES.items():
        uc.mem_write(addr, struct.pack('<I', val))

    def on_mem_read(uc_inst, access, addr, size, value, data):
        is_mmio = any(lo <= addr < hi for lo, hi, _ in MMIO_RANGES)
        if not is_mmio:
            return
        
        pc = last_pc[0]
        key = (pc, addr)
        poll_counts[key] = poll_counts.get(key, 0) + 1
        
        # If this PC+addr has been read more than 3 times, it's a poll
        # Inject the ready value to make it pass
        if poll_counts[key] == 3 and addr in READY_RESPONSES:
            val = READY_RESPONSES[addr]
            uc_inst.mem_write(addr, struct.pack('<I', val))
            injected[0] += 1
        
        # Also try: if any read from an unknown PHY offset, return
        # a value with common "ready" bits set
        if poll_counts[key] >= 3:
            off = addr & 0xFFFF
            if addr >= 0xFE0C0000 and addr < 0xFE100000:
                # PHY register - return with common ready bits
                uc_inst.mem_write(addr, struct.pack('<I', 0x00000003))
                injected[0] += 1
        
        val = struct.unpack_from('<I', bytes(uc_inst.mem_read(addr, 4)))[0]
        entry = {
            'i': count[0], 'op': 'R', 'addr': f'0x{addr:08X}',
            'reg': rname(addr), 'val': f'0x{val:08X}', 'pc': f'0x{pc:05X}',
            'poll': poll_counts[key]
        }
        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 not is_mmio:
            return
        pc = last_pc[0]
        entry = {
            'i': count[0], 'op': 'W', 'addr': f'0x{addr:08X}',
            'reg': rname(addr), 'val': f'0x{value:08X}', 'pc': f'0x{pc:05X}',
            'poll': 0
        }
        trace.append(entry)

    def on_code(uc_inst, addr, size, data):
        count[0] += 1
        last_pc[0] = addr
        unique_pcs.add(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)

    print(f"Running deep trace ({max_instr} instruction limit)...")
    try:
        uc.emu_start(0x10978, len(blob), timeout=120*1000000)
    except Exception as e:
        print(f"Exception: {e}")

    pc = uc.reg_read(UC_ARM64_REG_PC)
    x0 = uc.reg_read(UC_ARM64_REG_X0)
    
    print(f"\nExecution: {count[0]} instructions, {len(unique_pcs)} unique PCs")
    print(f"Final: PC=0x{pc:x} X0=0x{x0:x}")
    print(f"MMIO: {len(trace)} accesses ({injected[0]} ready values injected)")
    
    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)}")
    
    # Register access summary
    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 ({len(reg_counts)} unique registers):")
    print(f"{'Register':<35s} {'Reads':>6s} {'Writes':>7s}")
    print("-" * 50)
    for reg in sorted(reg_counts.keys()):
        c = reg_counts[reg]
        total = c['R'] + c['W']
        if total >= 2:  # skip single-access registers for brevity
            print(f"{reg:<35s} {c['R']:>6d} {c['W']:>7d}")
    
    # Execution phases (group accesses by PC range)
    print(f"\nExecution phases:")
    phases = {}
    for t in trace:
        pc_val = int(t['pc'], 16)
        phase = f"0x{(pc_val >> 12) << 12:05X}"
        if phase not in phases:
            phases[phase] = 0
        phases[phase] += 1
    for phase in sorted(phases.keys()):
        print(f"  {phase}: {phases[phase]} MMIO accesses")
    
    # Write trace
    if outfile:
        with open(outfile, 'w') as f:
            f.write("instr,op,addr,register,value,pc,poll_count\n")
            for t in trace:
                f.write(f"{t['i']},{t['op']},{t['addr']},{t['reg']},{t['val']},{t['pc']},{t['poll']}\n")
        print(f"\nTrace: {outfile}")

    # Show key moments
    print(f"\nKey MMIO sequence (first occurrence of each register):")
    seen = set()
    for t in trace:
        r = t['reg']
        if r not in seen:
            seen.add(r)
            print(f"  [{t['i']:7d}] {t['op']} {r:<30s} = {t['val']}  (PC={t['pc']})")

    return trace

if __name__ == '__main__':
    blob = sys.argv[1] if len(sys.argv) > 1 else '/opt/rkbin/bin/rk35/rk3588_ddr_lp4_2112MHz_lp5_2400MHz_v1.19.bin'
    out = sys.argv[2] if len(sys.argv) > 2 else '/opt/work/mmio_trace_deep.csv'
    run_deep_trace(blob, max_instr=2000000, outfile=out)