#!/usr/bin/env python3
"""
Patch MCU hex file to bypass CAN environment detection.

Three patches applied:
  1. 0x125B4: CAN/NM init function → return success immediately
     Original: B5FE 0004 (push; movs r4,r0) → 2001 4770 (movs r0,#1; bx lr)
  
  2. 0x124E8: Caller's NULL-pointer skip after CAN_init
     Original: B118 (cbz r0, 0x124F2) → BF00 (nop)
     This prevents the caller from skipping cleanup when CAN ptr is NULL
     
  3. 0x125D4: CAN search fail check → always continue
     Original: B1B0 (cbz r0, 0x12604) → BF00 (nop)
     This allows init to continue even if CAN hardware search fails
"""

import sys

HEX_FILE = '/Users/admin/workspace/jyq/mcubin/mcu_app.hex'
PATCHED_FILE = '/Users/admin/workspace/jyq/update/mcu_app_patched.hex'

PATCHES = [
    # (addr, original_bytes, new_bytes, description)
    (0x125B4, bytes([0xFE, 0xB5, 0x04, 0x00]), bytes([0x20, 0x01, 0x70, 0x47]),
     "CAN init: movs r0,#1; bx lr (return success)"),
    (0x124E8, bytes([0x18, 0xB1]), bytes([0x00, 0xBF]),
     "Caller NULL check: cbz -> nop (don't skip cleanup)"),
    (0x125D4, bytes([0xB0, 0xB1]), bytes([0x00, 0xBF]),
     "CAN search fail check: cbz -> nop (always continue)"),
    (0x125DE, bytes([0x1A, 0xD0]), bytes([0x00, 0xBF]),
     "CAN init fail check: beq -> nop (always continue)"),
    (0x125FC, bytes([0x17, 0xD0]), bytes([0x00, 0xBF]),
     "CAN final result check: beq -> nop (always continue)"),
]

def parse_hex(lines):
    """Parse Intel HEX into a dict of {addr: bytearray}"""
    data = {}
    extended_addr = 0
    for line in lines:
        line = line.strip()
        if not line or line[0] != ':':
            continue
        byte_count = int(line[1:3], 16)
        address = int(line[3:7], 16)
        record_type = int(line[7:9], 16)
        payload = bytes.fromhex(line[9:9+byte_count*2])
        
        if record_type == 4:
            extended_addr = int.from_bytes(payload, 'big') << 16
        elif record_type == 0:
            addr = extended_addr + address
            data[addr] = bytearray(payload)
    return data

def rebuild_hex(data, original_lines):
    """Rebuild Intel HEX, updating patched records with correct checksum"""
    new_lines = []
    extended_addr = 0
    
    for line in original_lines:
        line = line.strip()
        if not line or line[0] != ':':
            new_lines.append(line)
            continue
            
        byte_count = int(line[1:3], 16)
        address = int(line[3:7], 16)
        record_type = int(line[7:9], 16)
        
        if record_type == 4:
            payload = bytes.fromhex(line[9:9+byte_count*2])
            extended_addr = int.from_bytes(payload, 'big') << 16
            new_lines.append(line)
        elif record_type == 0:
            addr = extended_addr + address
            if addr in data:
                payload = bytes(data[addr])
                # Compute checksum
                csum = byte_count
                csum += (address >> 8) & 0xFF
                csum += address & 0xFF
                csum += record_type
                for b in payload:
                    csum += b
                csum = ((~csum) + 1) & 0xFF
                new_line = f":{byte_count:02X}{address:04X}00{payload.hex().upper()}{csum:02X}"
                new_lines.append(new_line)
            else:
                new_lines.append(line)
        else:
            new_lines.append(line)
    
    return new_lines

def main():
    with open(HEX_FILE, 'r') as f:
        lines = f.readlines()
    
    print(f"Read {len(lines)} lines from {HEX_FILE}")
    data = parse_hex(lines)
    print(f"Parsed {len(data)} data records")
    print()
    
    applied = []
    for addr, orig, new, desc in PATCHES:
        found = False
        for rec_addr, rec_data in data.items():
            if rec_addr <= addr < rec_addr + len(rec_data):
                offset = addr - rec_addr
                actual = bytes(rec_data[offset:offset+len(orig)])
                if actual == orig:
                    # Apply patch
                    for i, b in enumerate(new):
                        rec_data[offset + i] = b
                    print(f"  [OK] 0x{addr:05X}: {actual.hex()} -> {new.hex()}  ({desc})")
                    applied.append((addr, orig, new, desc))
                    found = True
                elif actual == new:
                    print(f"  [SKIP] 0x{addr:05X}: already patched ({desc})")
                    found = True
                else:
                    print(f"  [WARN] 0x{addr:05X}: expected {orig.hex()}, got {actual.hex()} ({desc})")
                    found = True
                break
        if not found:
            print(f"  [FAIL] 0x{addr:05X}: address not found in hex data ({desc})")
    
    if not applied:
        print("\nNo patches applied!")
        sys.exit(1)
    
    # Rebuild hex
    new_lines = rebuild_hex(data, lines)
    
    with open(PATCHED_FILE, 'w') as f:
        for line in new_lines:
            f.write(line + '\n')
    
    print(f"\nWrote {len(new_lines)} lines to {PATCHED_FILE}")
    print(f"\n=== {len(applied)} patches applied ===")
    for addr, orig, new, desc in applied:
        print(f"  0x{addr:05X}: {desc}")

if __name__ == '__main__':
    main()