#!/usr/bin/env python3
"""Patch original MCU hex file IN-PLACE, preserving record structure."""
import sys

HEX_FILE = '/Users/admin/workspace/jyq/mcubin/mcu_app.hex'
OUTPUT = '/Users/admin/workspace/jyq/update/mcu_app_patched_inplace.hex'

PATCHES = [
    (0x125B4, 'FEB50400', '20017047', 'CAN init: movs r0,#1; bx lr'),
    (0x124E8, '18B1', '00BF', 'Caller NULL check: cbz->nop'),
    (0x125D4, 'B0B1', '00BF', 'Search fail: cbz->nop'),
    (0x125DE, '1AD0', '00BF', 'Init fail: beq->nop'),
    (0x125FC, '17D0', '00BF', 'Final fail: beq->nop'),
]

def main():
    with open(HEX_FILE, 'r') as f:
        lines = f.readlines()
    print(f"Read {len(lines)} lines")

    extended = 0
    modified = 0
    applied = set()

    for i, line in enumerate(lines):
        line = line.strip()
        if not line or line[0] != ':':
            continue
        bc = int(line[1:3], 16)
        addr = int(line[3:7], 16)
        rt = int(line[7:9], 16)
        
        if rt == 4:
            extended = int(line[9:13], 16) << 16
            continue
        if rt != 0:
            continue
        
        start = extended + addr
        data = bytes.fromhex(line[9:9+bc*2])
        new_data = bytearray(data)
        any_change = False

        for p_addr, orig_hex, new_hex, desc in PATCHES:
            if p_addr in applied:
                continue
            orig = bytes.fromhex(orig_hex)
            new_b = bytes.fromhex(new_hex)
            if start <= p_addr < start + bc:
                offset = p_addr - start
                if data[offset:offset+len(orig)] == orig:
                    for j, b in enumerate(new_b):
                        new_data[offset + j] = b
                    any_change = True
                    applied.add(p_addr)
                    print(f"  [OK] 0x{p_addr:05X}: {orig_hex} -> {new_hex} ({desc})")
                else:
                    actual = data[offset:offset+len(orig)].hex()
                    print(f"  [WARN] 0x{p_addr:05X}: expected {orig_hex}, got {actual}")

        if any_change:
            payload_hex = new_data.hex().upper()
            csum = bc + (addr >> 8) + (addr & 0xFF) + rt
            for b in new_data:
                csum += b
            csum = ((~csum) + 1) & 0xFF
            lines[i] = f':{bc:02X}{addr:04X}{rt:02X}{payload_hex}{csum:02X}\n'
            modified += 1

    print(f"\nModified {modified} lines")
    print(f"Applied {len(applied)} of {len(PATCHES)} patches")

    if len(applied) != len(PATCHES):
        missing = [d for a,_,_,d in PATCHES if a not in applied]
        print(f"NOT applied: {missing}")
        sys.exit(1)

    with open(OUTPUT, 'w') as f:
        f.writelines(lines)
    print(f"Wrote {OUTPUT}")

if __name__ == '__main__':
    main()