#!/usr/bin/env python3
"""Cross-reference helper for the ESP8266 stock firmware (Xtensa LX106, call0 ABI).

Linear disassembly of the image goes out of sync on inline literal pools, so xrefs are
found by brute force: every byte offset is decoded as l32r / call0 and kept if it hits
the target. Results are then shown with objdump from a given start address, with each
l32r annotated with the literal value it loads.

    python tools/xref.py <dump.bin> lit 0x60000304      # literal slots holding a value + their l32r users
    python tools/xref.py <dump.bin> calls 0x40100abc    # call0 sites + l32r/callx users of a function
    python tools/xref.py <dump.bin> str "Light"         # string address(es) + l32r users
    python tools/xref.py <dump.bin> dis 0x40100abc 0x40100b20

Only the app image (0x1000 slot) is read. The objdump path defaults to the PlatformIO
toolchain; override with $OBJDUMP.
"""
import os
import re
import struct
import subprocess
import sys
import tempfile

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from esp8266_image_to_elf import APP_SLOT_MAX, build_elf, read_segments  # noqa: E402

OBJDUMP = os.environ.get("OBJDUMP") or os.path.expanduser(
    "~/.platformio/packages/toolchain-xtensa/bin/xtensa-lx106-elf-objdump")


class Image:
    def __init__(self, dump, offset=0x1000):
        with open(dump, "rb") as f:
            f.seek(offset)
            raw = f.read(APP_SLOT_MAX - offset)
        self.entry, self.segs = read_segments(raw)
        self._elf = None

    def code_segs(self):
        return [(a, d) for a, d in self.segs if a >= 0x40000000]

    def seg_of(self, addr):
        for a, d in self.segs:
            if a <= addr < a + len(d):
                return a, d
        return None

    def word(self, addr):
        s = self.seg_of(addr)
        if not s:
            return None
        a, d = s
        o = addr - a
        return struct.unpack_from("<I", d, o)[0] if o + 4 <= len(d) else None

    def literals(self, value):
        """Aligned words equal to value (literal pool slots)."""
        out = []
        pat = struct.pack("<I", value)
        for a, d in self.segs:
            i = d.find(pat)
            while i >= 0:
                if (a + i) % 4 == 0:
                    out.append(a + i)
                i = d.find(pat, i + 1)
        return out

    def l32r_users(self, lit_addr):
        """(pc, reg) of every byte offset that decodes as l32r loading lit_addr."""
        out = []
        for a, d in self.code_segs():
            for i in range(len(d) - 2):
                b0 = d[i]
                if b0 & 0xF != 1:
                    continue
                imm16 = d[i + 1] | (d[i + 2] << 8)
                pc = a + i
                tgt = (((pc + 3) & ~3) + ((imm16 | ~0xFFFF) << 2)) & 0xFFFFFFFF
                if tgt == lit_addr:
                    out.append((pc, b0 >> 4))
        return out

    def call0_sites(self, func):
        """Byte offsets that decode as call0 to func (CALL format: op0=5, n=0)."""
        out = []
        for a, d in self.code_segs():
            for i in range(len(d) - 2):
                v = d[i] | (d[i + 1] << 8) | (d[i + 2] << 16)
                if v & 0x3F != 0x05:
                    continue
                off = v >> 6
                if off & 0x20000:
                    off -= 0x40000
                pc = a + i
                if ((pc & ~3) + (off << 2) + 4) & 0xFFFFFFFF == func:
                    out.append(pc)
        return out

    def elf(self):
        if not self._elf:
            fd, path = tempfile.mkstemp(suffix=".elf")
            with os.fdopen(fd, "wb") as f:
                f.write(build_elf(self.entry, self.segs))
            self._elf = path
        return self._elf

    def dis(self, start, stop):
        out = subprocess.run([OBJDUMP, "-d", f"--start-address={start}",
                              f"--stop-address={stop}", self.elf()],
                             capture_output=True, text=True).stdout
        lines = []
        for ln in out.splitlines():
            m = re.match(r"\s*([0-9a-f]+):\s", ln)
            if not m:
                continue
            m2 = re.search(r"l32r\s+a\d+, 0x([0-9a-f]+)", ln)
            if m2:
                w = self.word(int(m2.group(1), 16))
                if w is not None:
                    ln += f"    ; ={w:#010x}{self.describe(w)}"
            lines.append(ln)
        return "\n".join(lines)

    def ctx(self, addr, before=24, after=12):
        """Disassembly around addr, starting at an offset that decodes in step with addr."""
        tag = f"{addr:x}:"
        for back in range(before, before + 16):
            text = self.dis(addr - back, addr + after)
            lines = text.splitlines()
            if any(ln.strip().startswith(tag) for ln in lines):
                return "\n".join((">" if ln.strip().startswith(tag) else " ") + ln for ln in lines)
        return self.dis(addr - before, addr + after)

    def describe(self, w):
        """Short hint for literal values: a printable string or a known register."""
        regs = {0x60000300: "GPO", 0x60000304: "GPOS", 0x60000308: "GPOC",
                0x6000030c: "GPE", 0x60000310: "GPES", 0x60000314: "GPEC",
                0x60000318: "GPI", 0x6000031c: "GPIE", 0x60000320: "GPIES",
                0x60000324: "GPIEC", 0x60000800: "IOMUX", 0x60000768: "GP16O",
                0x60000774: "GP16E", 0x6000078c: "GP16I", 0x60000790: "GPF16",
                0x60000100: "SPI1CMD", 0x60000140: "SPI1W0"}
        if w in regs:
            return f" {regs[w]}"
        if 0x60000000 <= w < 0x60001000:
            return " (periph)"
        s = self.seg_of(w)
        if s and w < 0x40300000:
            a, d = s
            o = w - a
            end = d.find(b"\0", o, o + 64)
            if end > o + 2:
                t = d[o:end]
                if all(32 <= c < 127 for c in t):
                    return f' "{t.decode()}"'
        return ""


def main():
    if len(sys.argv) < 4:
        raise SystemExit(__doc__)
    img = Image(sys.argv[1])
    cmd, arg = sys.argv[2], sys.argv[3]
    if cmd == "lit":
        for lit in img.literals(int(arg, 0)):
            users = ", ".join(f"{pc:#x}(a{r})" for pc, r in img.l32r_users(lit))
            print(f"literal @ {lit:#x}: {users or '-'}")
    elif cmd == "calls":
        f = int(arg, 0)
        for pc in img.call0_sites(f):
            print(f"call0 @ {pc:#x}")
        for lit in img.literals(f):
            users = ", ".join(f"{pc:#x}(a{r})" for pc, r in img.l32r_users(lit))
            print(f"literal @ {lit:#x}: {users or '-'}")
    elif cmd == "str":
        pat = arg.encode()
        for a, d in img.segs:
            i = d.find(pat)
            while i >= 0:
                sa = a + i
                users = []
                for lit in img.literals(sa):
                    users += [f"{pc:#x}(a{r})" for pc, r in img.l32r_users(lit)]
                print(f"{sa:#x}: {', '.join(users) or '-'}")
                i = d.find(pat, i + 1)
    elif cmd == "dis":
        print(img.dis(int(arg, 0), int(sys.argv[4], 0)))
    elif cmd == "ctx":
        before = int(sys.argv[4], 0) if len(sys.argv) > 4 else 24
        after = int(sys.argv[5], 0) if len(sys.argv) > 5 else 12
        for a in arg.split(","):
            print(f"--- {a}")
            print(img.ctx(int(a, 0), before, after))
    else:
        raise SystemExit(__doc__)


if __name__ == "__main__":
    main()