import sys
import os


def patch_binary(file_path, offset, expected_bytes, patch_bytes):
    if not os.path.exists(file_path):
        print(f"Error: File {file_path} not found.")
        return False

    with open(file_path, "rb") as f:
        data = bytearray(f.read())

    if offset + len(expected_bytes) > len(data):
        print(f"Error: Offset {hex(offset)} is out of range.")
        return False

    current_bytes = data[offset : offset + len(expected_bytes)]
    if current_bytes != expected_bytes:
        print(f"Error: Bytes at {hex(offset)} mismatch.")
        print(f"Expected: {expected_bytes.hex()}")
        print(f"Found:    {current_bytes.hex()}")
        return False

    print(f"Patching {file_path} at {hex(offset)}...")
    data[offset : offset + len(patch_bytes)] = patch_bytes

    output_path = file_path.replace(".bin", ".patched.bin")
    with open(output_path, "wb") as f:
        f.write(data)

    print(f"Successfully wrote patched binary to {output_path}")
    return True


if __name__ == "__main__":
    ultra_original = bytes.fromhex("26153c")
    ultra_patch = bytes.fromhex("060f00")

    patch_binary(
        "/home/sapient/Public/03_Projects_and_Dev/Meshcore_TDECK/Ultra-TDeck-v9.2.bin",
        0xBA62D,
        ultra_original,
        ultra_patch,
    )
    patch_binary(
        "/home/sapient/Public/esp32-analysis/firmware/Ultra-TDeck-v9.2-merged.bin",
        0xCA62D,
        ultra_original,
        ultra_patch,
    )

    meshos_original = bytes.fromhex("26193c")
    meshos_patch = bytes.fromhex("060f00")

    patch_binary(
        "/home/sapient/Public/03_Projects_and_Dev/Meshcore_TDECK/MeshOS-TDeck-1.1.8.bin",
        0xB99ED,
        meshos_original,
        meshos_patch,
    )

    merged_path = (
        "/home/sapient/Public/esp32-analysis/firmware/MeshOS-TDeck-1.1.8-merged.bin"
    )
    if os.path.exists(merged_path):
        patch_binary(merged_path, 0xC99ED, meshos_original, meshos_patch)