#!/usr/bin/python

import sys, zlib

class Injector:

    vacancies = [
        (0x00ed83a0, 0x01d5ca00), # Normal files spot
        (0x0002ea70, 0x00039850), # Unused space 1
        (0x00157810, 0x001a15c0), # Unused space 2
    ]

    def rarezip(self, buffer):
        length = len(buffer)

        header = bytes([0x11, 0x73])
        lendata = bytes([length >> 16])
        lendata += bytes([(length >> 8) & 0xff])
        lendata += bytes([length & 0xff])

        obj = zlib.compressobj(level=9, wbits=-15)
        return header + lendata + obj.compress(buffer) + obj.flush()

    def rareunzip(self, compressed):
        return zlib.decompress(compressed[5:], wbits=-15)

    def load(self, romfile):
        fp = open(romfile, 'rb')
        rombuffer = fp.read()
        fp.close()

        self.rompart1 = rombuffer[0:0x2ea70]
        self.rompart2 = rombuffer[0x0004e850:0x00157810]
        self.rompart3 = rombuffer[0x001a15c0:0x00ed83a0]
        self.rompart4 = rombuffer[0x01d5ca00:]

        fp = open('build/ntsc-final/ucode/setup.bin', 'rb')
        setup = fp.read()
        fp.close()

        self.globaltop = setup[0:0x28080]
        self.globalbot = setup[0x2a000:]

        self.files = []
        i = 0

        while i <= 0x7dd:
            romnameaddr = int.from_bytes(rombuffer[0x01d5ca00 + i * 4:0x01d5ca00 + i * 4 + 4], 'big') + 0x01d5ca00
            romdataaddr = int.from_bytes(setup[0x28080 + i * 4:0x28080 + i * 4 + 4], 'big')

            name = ''
            while rombuffer[romnameaddr] != 0:
                name = name + chr(rombuffer[romnameaddr])
                romnameaddr += 1

            self.files.append({
                'name': name,
                'romaddr': romdataaddr,
                'data': None,
            })
            i += 1

        self.files.sort(key=lambda file: file['romaddr'])

        for index, file in enumerate(self.files):
            start = file['romaddr']
            end = self.files[index + 1]['romaddr'] if index < 0x7dd else 0x01d5ca00
            file['data'] = rombuffer[start:end]

    def transform(self):
        # Replace file contents
        for file in self.files:
            if self.isFilePointless(file['name']):
                file['data'] = None
                continue

            # Zipped file
            try:
                fp = open('build/ntsc-final/files/%s' % file['name'], 'rb')
                contents = fp.read()
                fp.close()
            except:
                continue

            file['data'] = self.align(contents)

        # Calculate new ROM addresses
        vacancyid = 0
        romaddr = self.vacancies[0][0]
        vacend = self.vacancies[0][1]

        for file in self.files:
            if file['name'] == '':
                continue

            filelen = len(file['data'])
            available = vacend - romaddr

            if filelen > available:
                vacancyid += 1
                romaddr = self.vacancies[vacancyid][0]
                vacend = self.vacancies[vacancyid][1]
                available = vacend - romaddr

            file['romaddr'] = romaddr
            romaddr += filelen

    def isFilePointless(self, name):
        return False

    def align(self, buffer):
        length = len(buffer)

        if length % 0x10 == 0:
            return buffer

        over = length % 0x10
        pad = 0x10 - over

        buffer += (0).to_bytes(1, 'big') * pad

        return buffer

    def compile(self):
        buffer = self.rompart1
        assert(len(buffer) == 0x2ea70)
        buffer += self.compileVacancy(1)
        assert(len(buffer) == 0x39850)
        buffer += self.compileGlobals()
        assert(len(buffer) == 0x4e850)
        buffer += self.rompart2
        assert(len(buffer) == 0x157810)
        buffer += self.compileVacancy(2)
        assert(len(buffer) == 0x1a15c0)
        buffer += self.rompart3
        assert(len(buffer) == 0xed83a0)
        buffer += self.compileVacancy(0)
        assert(len(buffer) == 0x01d5ca00)
        buffer += self.rompart4
        assert(len(buffer) == 0x2000000)
        return buffer

    def compileGlobals(self):
        buffer = self.globaltop

        for file in self.files:
            buffer += file['romaddr'].to_bytes(4, 'big')

        buffer += (0x01d5ca00).to_bytes(4, 'big')
        buffer += (0).to_bytes(4, 'big')
        buffer += self.globalbot

        buffer = self.rarezip(buffer)

        available = 0x15000 - len(buffer)
        buffer += (0).to_bytes(1, 'big') * available
        return buffer

    def compileVacancy(self, vacid):
        buffer = bytes()
        vacstart = self.vacancies[vacid][0]
        vacend = self.vacancies[vacid][1]

        for file in self.files:
            if file['romaddr'] >= vacstart and file['romaddr'] < vacend:
                buffer += file['data']

        available = vacend - vacstart - len(buffer)
        buffer += (0).to_bytes(1, 'big') * available

        return buffer

    def ROL(self, i, b):
        return ((i << b) | (i >> (32 - b))) & 0xffffffff

    def R4(self, b):
        return b[0]*0x1000000 + b[1]*0x10000 + b[2]*0x100 + b[3]

    def crc(self, f):
        seed = 0xdf26f436
        t1 = t2 = t3 = t4 = t5 = t6 = seed

        f.seek(0x0710 + 0x40)
        lookup = f.read(0x100)

        f.seek(0x1000)
        for i in range(0x1000, 0x101000, 4):
            d = self.R4(f.read(4))

            if ((t6 + d) & 0xffffffff) < t6:
                t4 += 1
                t4 &= 0xffffffff

            t6 += d
            t6 &= 0xffffffff

            t3 ^= d

            r = self.ROL(d, d & 0x1F)

            t5 += r
            t5 &= 0xffffffff

            if t2 > d:
                t2 ^= r
            else:
                t2 ^= t6 ^ d

            o = i & 0xFF
            temp = self.R4(lookup[o:o + 4])
            t1 += temp ^ d
            t1 &= 0xffffffff

        crc1 = t6 ^ t4 ^ t3
        crc2 = t5 ^ t2 ^ t1

        return crc1 & 0xffffffff, crc2 & 0xffffffff

injector = Injector()
injector.load(sys.argv[1])
injector.transform()

filename = 'build/ntsc-final/pd.z64'

fp = open(filename, 'wb')
fp.write(injector.compile())
fp.close()

fp = open(filename, 'rb')
crcs = injector.crc(fp)
fp.seek(0)
buffer = fp.read()
fp.close()

fp = open(filename, 'r+b')
fp.seek(0x10)
fp.write(crcs[0].to_bytes(4, 'big'))
fp.write(crcs[1].to_bytes(4, 'big'))
fp.close()