#!/usr/bin/env python3
# Boktai 1 checksum fixer
# Usage: python3 boktai1_checksum_fix.py SAVEGAME.SaveRAM

import struct
import sys

# Memory blocks that are checksummed, array of (start, length) pairs, in bytes
BLOCKS = [
    (0x0000, 0x18),
    (0x00c0, 0x28),
    (0x0200, 0x28),
    (0x0340, 0x400),
    (0x0780, 0x418),
    (0x0dc0, 0x400),
    (0x1200, 0x418),
    (0x1840, 0x168),
    (0x1a40, 0x168),
]

# Calculates checksum of xs, an array of uint16_t's. Returns 4 uint16_t's.
def checksum(xs):
    ck_add = 0
    ck_xor = 0
    for x in xs:
        ck_add = (ck_add + x) & 0xffff
        ck_xor ^= x

    return [
        # Constant
        0xffff,
        # Number of bytes in this block
        len(xs)*2,
        # All halfwords folded with addition
        ck_add,
        # All halfwords folded with XOR
        ck_xor,
    ]

if len(sys.argv) < 2:
    print("Usage: python3 " + sys.argv[0] + " SAVEGAME.SaveRAM")
    sys.exit(1)

with open(sys.argv[1], "r+b") as f:
    for start, length in BLOCKS:
        f.seek(start)
        block = f.read(length)
        # The save is an array of big-endian halfwords, so we use > to convert
        # to native endianess when reading and writing
        halfwords = struct.iter_unpack(">H", block)
        check = checksum([x[0] for x in halfwords])
        for c in check:
            f.write(struct.pack(">H", c))
print("OK")