#!/usr/bin/env python3
"""verify_patch_asm.py -- prove the sdas8051-assembled patch matches the
bytes actually written to the device by patch_c2.py.

Pipeline:
  1. Parse patch_usb1_to_cv.lst (sdas8051 listing) -> {addr: bytes}.
  2. Extract the assembled stub (0x8126, 23 B) and retarget (0xDF28, 3 B).
  3. Compare against patch_usb1_to_cv.STUB / NEW_CALL (the authoritative
     bytes patch_c2.py verified on the device over C2).
  4. Splice the assembled bytes into the stock binary and diff the result
     against the expected patched image (stock + STUB + NEW_CALL).

Run from this directory:
  sdas8051 -l -s patch_usb1_to_cv.asm
  python3 verify_patch_asm.py
"""
import re
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
FWTOOLS = HERE.parent
sys.path.insert(0, str(FWTOOLS))
import patch_usb1_to_cv as p  # noqa: E402

BASE = 0x2400  # app image load address (flat binary base)


def parse_lst_bytes(lst_path):
    """Return {addr: bytes} for every instruction/data line in the listing."""
    out = {}
    # "      008126 12 A5 7B         [24]   30     lcall ..."
    line_re = re.compile(r'^\s+([0-9a-fA-F]{6})\s+([0-9a-fA-F]{2}(?:\s+[0-9a-fA-F]{2})*)\s+\[\d+\]')
    for ln in Path(lst_path).read_text(errors="replace").splitlines():
        m = line_re.match(ln)
        if not m:
            continue
        addr = int(m.group(1), 16)
        bs = bytes(int(b, 16) for b in m.group(2).split())
        out[addr] = bs
    return out


def splice(stock, addr, data):
    """Overlay `data` at file offset (addr - BASE) in a copy of stock."""
    off = addr - BASE
    img = bytearray(stock)
    img[off:off + len(data)] = data
    return bytes(img)


def main():
    lst = HERE / "patch_usb1_to_cv.lst"
    if not lst.exists():
        sys.exit("missing patch_usb1_to_cv.lst -- run: sdas8051 -l -s patch_usb1_to_cv.asm")
    mem = parse_lst_bytes(lst)

    # Assembled patch bytes (contiguous).
    asm_stub = mem[p.STUB_ADDR]
    asm_call = mem[p.CALL_SITE]
    # Sanity: stub must be exactly 23 contiguous bytes 0x8126..0x813C.
    expected_stub_len = len(p.STUB)
    assert len(asm_stub) == 3, asm_stub  # first instruction is lcall (3 B); listing gives per-line bytes
    # Reassemble the full stub by walking consecutive listing entries.
    stub_bytes = b""
    a = p.STUB_ADDR
    while len(stub_bytes) < expected_stub_len:
        chunk = mem[a]
        stub_bytes += chunk
        a += len(chunk)
    assert len(stub_bytes) == expected_stub_len, (len(stub_bytes), expected_stub_len)

    print("assembled stub :", stub_bytes.hex(' '))
    print("expected  STUB :", p.STUB.hex(' '))
    print("assembled call :", asm_call.hex(' '))
    print("expected  CALL :", p.NEW_CALL.hex(' '))

    ok_stub = stub_bytes == bytes(p.STUB)
    ok_call = asm_call == bytes(p.NEW_CALL)
    print(f"\nstub match : {ok_stub}")
    print(f"call match : {ok_call}")

    # Full-image diff: stock + assembled patch  vs  stock + authoritative patch.
    stock = (HERE / "QuNexus_Firmware_v2.2.1.bin").read_bytes()
    img_asm = splice(stock, p.STUB_ADDR, stub_bytes)
    img_asm = splice(img_asm, p.CALL_SITE, asm_call)
    img_exp = splice(stock, p.STUB_ADDR, bytes(p.STUB))
    img_exp = splice(img_exp, p.CALL_SITE, bytes(p.NEW_CALL))

    print(f"\nstock size      : {len(stock)} bytes")
    print(f"patched (asm)   : {len(img_asm)} bytes")
    print(f"patched (expect): {len(img_exp)} bytes")
    print(f"full-image match: {img_asm == img_exp}")

    if ok_stub and ok_call and img_asm == img_exp:
        print("\nPASS: sdas8051-assembled patch is byte-identical to the patch "
              "written to the device.")
        return 0
    print("\nFAIL")
    return 1


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