#!/usr/bin/env python3
"""RT-4D firmware xref/disasm helper.

  python xref.py dis  <vaddr> [n]      disassemble n bytes at vaddr
  python xref.py callers <vaddr>       who bl/blx's this address
  python xref.py lit  <value>          where this 32-bit literal appears (pool refs)
  python xref.py imm  <value>          movw/movt pairs building this constant
  python xref.py find <hexbytes>       raw byte search
"""
import sys, struct, capstone

BASE = 0x08002800
import os
_CANDS = [
    "/home/viktor/claude/rt-4d/stock-fw/rt4d_stock_v3.25_abs_0x08002800.bin",
    "/home/viktor/claude/rt-4d-repo/firmware/rt4d_stock_v3.25_abs_0x08002800.bin",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "firmware",
                 "rt4d_stock_v3.25_abs_0x08002800.bin"),
]
IMG = next(open(p, "rb").read() for p in _CANDS if os.path.exists(p))

md = capstone.Cs(capstone.CS_ARCH_ARM, capstone.CS_MODE_THUMB)
md.detail = True


def _word(va):
    o = va - BASE
    if 0 <= o <= len(IMG) - 4:
        return struct.unpack_from("<I", IMG, o)[0]
    return None


def dis(va, n=0x60):
    """disassemble, annotating PC-relative literal loads with the resolved value"""
    off = va - BASE
    import re as _re
    for i in md.disasm(IMG[off:off + n], va):
        ann = ""
        m = _re.search(r"\[pc, #(?:0x)?([0-9a-fA-F]+)\]", i.op_str)
        if m and i.mnemonic.startswith("ldr"):
            pcv = ((i.address + 4) & ~3) + int(m.group(1), 16)
            w = _word(pcv)
            if w is not None:
                ann = f"   ; @0x{pcv:08X} = 0x{w:08X}"
        print(f"  0x{i.address:08X}: {i.bytes.hex():<10} {i.mnemonic:8} {i.op_str}{ann}")


def _all_calls():
    """Linear sweep collecting bl/blx targets -> {target: [sites]}."""
    xr = {}
    for start in (0, 2):                       # both alignments, dedupe by site
        for i in md.disasm(IMG[start:], BASE + start):
            if i.mnemonic in ("bl", "blx") and i.op_str.startswith("#"):
                t = int(i.op_str[1:], 16)
                xr.setdefault(t, set()).add(i.address)
    return {k: sorted(v) for k, v in xr.items()}


def callers(va):
    xr = _all_calls()
    hits = xr.get(va, []) + xr.get(va | 1, []) + xr.get(va & ~1, [])
    hits = sorted(set(hits))
    print(f"callers of 0x{va:08X}: {len(hits)}")
    for h in hits:
        print(f"  0x{h:08X}")
    return hits


def lit(val):
    hits = []
    for off in range(0, len(IMG) - 4, 4):
        if struct.unpack_from("<I", IMG, off)[0] == val:
            hits.append(BASE + off)
    print(f"literal 0x{val:08X}: {len(hits)} hit(s)")
    for h in hits[:60]:
        print(f"  0x{h:08X}")
    return hits


def imm(val):
    """find movw/movt pairs that build `val` (constant loaded into a reg)."""
    lo, hi = val & 0xFFFF, (val >> 16) & 0xFFFF
    out = []
    for i in md.disasm(IMG, BASE):
        if i.mnemonic == "movw" and i.op_str.endswith(f"#{hex(lo)}"):
            out.append(i.address)
    print(f"movw #{hex(lo)} (for 0x{val:08X}): {len(out)} site(s)")
    for a in out[:60]:
        print(f"  0x{a:08X}")
    return out


def find(hexs):
    pat = bytes.fromhex(hexs)
    off = 0
    hits = []
    while True:
        j = IMG.find(pat, off)
        if j < 0:
            break
        hits.append(BASE + j); off = j + 1
    print(f"bytes {hexs}: {len(hits)} hit(s)")
    for h in hits[:60]:
        print(f"  0x{h:08X}")
    return hits


if __name__ == "__main__":
    cmd = sys.argv[1]
    if cmd == "dis":
        dis(int(sys.argv[2], 16), int(sys.argv[3], 0) if len(sys.argv) > 3 else 0x60)
    elif cmd == "callers":
        callers(int(sys.argv[2], 16))
    elif cmd == "lit":
        lit(int(sys.argv[2], 16))
    elif cmd == "imm":
        imm(int(sys.argv[2], 16))
    elif cmd == "find":
        find(sys.argv[2])