#!/usr/bin/env python3

import os
import zlib

"""
mkgamezips - Creates the segments/gamezips.bin from segments/game.bin

game.bin is the compiled game code from ld. This game code is split into
4KB chunks. Each chunk is individually zipped.

The format of the gamezips binary is:
* Array of offsets to each zip, where each offset is 4 bytes and relative to the
  start of the gamezips segment.
* After the array of offsets comes the data it points to. Each data consists of:
    * A 2 byte checksum of the uncompressed chunk.
    * Zip data (starting with 0x1173001000).
    * Optional single byte to align it to the next 2 byte boundary. The added
      byte is probable garbage data.
"""

def main():
    bins = get_bins()
    sums = get_sums(bins)
    zips = get_zips(bins)

    fd = open('build/%s/segments/gamezips.bin' % os.environ['ROMID'], 'wb+')
    pos = len(zips) * 4 + 4

    # Write pointer array
    for zip in zips:
        fd.write(pos.to_bytes(4, byteorder='big'))
        pos += 2 + len(zip)
        if pos % 2 == 1:
            pos += 1

    # Last pointer points to end
    fd.write(pos.to_bytes(4, byteorder='big'))

    # Write checksums and compressed data
    padding_index = 0
    for index, zip in enumerate(zips):
        if pos % 2 == 1:
            try:
                pad_value = padding[os.environ['ROMID']][padding_index]
            except IndexError:
                pad_value = 0
            fd.write(pad_value.to_bytes(1, 'big'))

            padding_index += 1
            pos += 1

        fd.write(sums[index].to_bytes(2, 'big'))
        fd.write(zip)

        pos += len(zip)

    # The previous two bytes are repeated at the end
    fd.seek(-2, os.SEEK_CUR)
    value = fd.read(2)
    fd.write(value)
    fd.close()

def get_filecontents(filename):
    fd = open(filename, 'rb')
    binary = fd.read()
    fd.close()
    return binary

def get_bins():
    binary = get_filecontents('build/%s/segments/game.bin' % os.environ['ROMID'])
    return [binary[i:i+0x1000] for i in range(0, len(binary), 0x1000)]

def get_sums(bins):
    return [sum(part) for part in bins]

def get_zips(bins):
    return [zip(part) for part in bins]

def sum(binary):
    sum = 0
    for i in range(0, len(binary), 4):
        sum += int.from_bytes(binary[i:i+4], 'big')
    return sum & 0xffff

def zip(binary):
    obj = zlib.compressobj(level=9, wbits=-15)
    return b'\x11\x73' + len(binary).to_bytes(3, 'big') + obj.compress(binary) + obj.flush()

padding = {
    'ntsc-beta': [],
    'ntsc-1.0': [
        0x00, 0x00, 0x00, 0x73, 0x75, 0x47, 0xa4, 0x79,
        0x91, 0x1e, 0xd8, 0xd8, 0x7b, 0xe2, 0x13, 0x51,
        0xaa, 0x2a, 0xa1, 0x5a, 0xbd, 0x24, 0x87, 0x0d,
        0xdc, 0x50, 0xd3, 0xe2, 0x39, 0x9f, 0x1a, 0x86,
        0x73, 0xc5, 0x72, 0xdd, 0xc8, 0xa9, 0xd9, 0x21,
        0xf2, 0xd4, 0x66, 0xa5, 0xbb, 0xfa, 0xbe, 0x6d,
        0xc4, 0xce, 0x63, 0x6f, 0xba, 0xf7, 0xe3, 0xc9,
        0xd4, 0xd3, 0xfb, 0x9c, 0xff, 0x76, 0xcf, 0x20,
        0x6e, 0xbb, 0x2c, 0x90, 0xd3, 0xc1, 0xc4, 0x9f,
        0xb4, 0x70, 0x38, 0xe3, 0x22, 0x3e, 0xec, 0x5c,
        0xa4, 0xdc, 0x39, 0xe8, 0x8b, 0x5e, 0xcc, 0x3e,
        0xa2, 0x62, 0x87, 0x20, 0xd3, 0x05, 0x1c, 0xa5,
        0xd4, 0xca, 0x0e, 0xef, 0xed, 0x23, 0x67, 0xc6,
        0xff, 0x69, 0x66, 0x66, 0x9d, 0xaf, 0xb1, 0x2e,
        0x61, 0xc7, 0x9c, 0x56, 0x60, 0xe6, 0x95, 0xf8,
        0x21, 0x79, 0x93, 0x37, 0xe1, 0x8c, 0xb2, 0xcf,
        0x26, 0x01, 0x9d, 0xf8, 0x6a, 0x8b, 0x71, 0x79,
        0x57, 0xef, 0xaf, 0xb2, 0x27, 0xad, 0x4a, 0xf2,
        0xcc, 0x46, 0x3e, 0x34, 0x0f, 0x31, 0xec, 0xe3,
        0xc2, 0x11, 0xee, 0xb1, 0x4e, 0x3b, 0xcc, 0xd1,
        0xad, 0x4f, 0xbe, 0xef, 0x7b, 0x39, 0xd6, 0x97,
        0xf5, 0xf7, 0x02, 0xa8, 0x36, 0xa1, 0xd7, 0x50,
        0xd6, 0xaf, 0xdf, 0xb9, 0x22, 0x73, 0x4a, 0x37,
        0xf9, 0x46, 0x9d, 0x30, 0x1d, 0x1a, 0x13, 0x1c,
        0x9a, 0xa1, 0x22, 0xa7, 0xb4, 0x9b, 0x88, 0x19,
        0xd4, 0x00, 0xb9, 0xcf, 0xe6, 0xce, 0xf2, 0xd8,
        0xfd, 0x9c, 0x26, 0xe9, 0x4c, 0x59, 0x3d, 0xa6,
        0xa6, 0xa6, 0x07, 0x29, 0x24, 0x7a, 0xba,
    ],
    'ntsc-final': [
        0x00, 0xc2, 0xf9, 0x3a, 0x4d, 0x18, 0x8a, 0x07,
        0x4c, 0x68, 0x38, 0x17, 0x3c, 0x94, 0x98, 0x25,
        0x82, 0x6f, 0xf7, 0x2e, 0x41, 0xed, 0xe4, 0x88,
        0xf5, 0x48, 0x59, 0x5a, 0x7a, 0xb4, 0x3b, 0xf9,
        0xc7, 0xc7, 0xab, 0xf6, 0xef, 0x23, 0x8b, 0x6a,
        0x58, 0xb0, 0x27, 0x84, 0x77, 0x88, 0xe1, 0x34,
        0xf5, 0xf4, 0xa9, 0xb9, 0x30, 0x8a, 0x64, 0x23,
        0xb5, 0x6c, 0x87, 0xff, 0xd4, 0x84, 0xe4, 0x7c,
        0x93, 0xa0, 0x5b, 0x41, 0x28, 0xfa, 0x65, 0x3e,
        0xad, 0x51, 0x35, 0xf9, 0xec, 0x6a, 0xe9, 0xaf,
        0xe0, 0x7f, 0xe5, 0x8e, 0x0e, 0x6b, 0x42, 0x97,
        0xee, 0xad, 0x5d, 0xba, 0x91, 0x7c, 0xd6, 0x91,
        0xb3, 0xbd, 0x5f, 0x3d, 0x48, 0xd1, 0x37, 0xba,
        0xfc, 0x83, 0x51, 0x2b, 0xcb, 0x2f, 0x6b, 0xbf,
        0xb0, 0xe5, 0x9c, 0xac, 0x1d, 0x63, 0xcb, 0xa5,
        0x5e, 0x66, 0x24, 0x2d, 0xe3, 0x86, 0x3e, 0x0c,
        0xcf, 0x1a, 0x57, 0xc9, 0x4b, 0x29, 0x70, 0x31,
        0xbb, 0x55, 0xc4, 0x42, 0x62, 0x5e, 0x9a, 0xa0,
        0xff, 0x41, 0xf5, 0x62, 0x9f, 0xc9, 0x61, 0xee,
        0xbe, 0x7b, 0x0e, 0xf2, 0xd7, 0xf1, 0x90, 0x69,
        0xfa, 0xff, 0xf7, 0xb1, 0x3a, 0x27, 0xac, 0xc2,
        0x57, 0x7e, 0xcc, 0x92, 0xdd, 0x2d, 0x63, 0xa0,
        0x53, 0x74, 0x35, 0xbb, 0x24, 0xde, 0x6d, 0xbb,
        0x2c, 0xe5, 0xff, 0xeb, 0x37, 0xde, 0xd0, 0x6e,
        0x96, 0xfa, 0xbe, 0x79, 0xe3, 0x1e, 0x7f, 0xff,
        0x67, 0x86, 0x86, 0x86, 0x15, 0x6e,
    ],
    'pal-final': [],
}

main()