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


def apply_global_status_patch(file_path, offset, expected_bytes, patch_bytes):
    if not os.path.exists(file_path):
        return False

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

    if offset + len(expected_bytes) > len(data):
        return False

    current_bytes = data[offset : offset + len(expected_bytes)]
    if current_bytes != expected_bytes:
        print(
            f"Skipping Global Status patch for {file_path} at {hex(offset)}: mismatch."
        )
        print(f"Expected: {expected_bytes.hex()}")
        print(f"Found:    {current_bytes.hex()}")
        return False

    print(f"Applying Global Status patch to {file_path} at {hex(offset)}...")
    data[offset : offset + len(patch_bytes)] = patch_bytes

    # Overwrite the file since it's already a patched version or we want to update it
    with open(file_path, "wb") as f:
        f.write(data)

    return True


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

    ultra_global_original = bytes.fromhex("040242")
    ultra_global_patch = bytes.fromhex("221000")

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

    # Also apply global status patch to Ultra
    apply_global_status_patch(
        "/home/sapient/Public/03_Projects_and_Dev/Meshcore_TDECK/Ultra-TDeck-v9.2.patched.bin",
        0xBA6AD,
        ultra_global_original,
        ultra_global_patch,
    )

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

    apply_global_status_patch(
        "/home/sapient/Public/esp32-analysis/firmware/Ultra-TDeck-v9.2-merged.patched.bin",
        0xCA6AD,
        ultra_global_original,
        ultra_global_patch,
    )

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

    meshos_global_original = bytes.fromhex("84210040")
    meshos_global_patch = bytes.fromhex("221000")

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

    apply_global_status_patch(
        "/home/sapient/Public/03_Projects_and_Dev/Meshcore_TDECK/MeshOS-TDeck-1.1.8.patched.bin",
        0xB9A54,
        meshos_global_original,
        meshos_global_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)
        apply_global_status_patch(
            merged_path.replace(".bin", ".patched.bin"),
            0xC9A54,
            meshos_global_original,
            meshos_global_patch,
        )