#!/usr/bin/env python3

import os
import re
import subprocess

def main():
    fd = open(bdir() + '/pd.z64', 'wb+')

    # The retail ROM contains truncated duplicates of some segments.
    # For example, the real boot segment is at 0x1000 - 0x3050, but the tail end
    # of it is repeated at 0x2ea6c - 0x30a60. The truncated parts are not read
    # by the ROM; they are likely a side effect of Rare's linker copying things
    # around in the ROM.
    if os.environ['ROMID'] == 'ntsc-final':
        write_binary(fd, 0x2ea1c, get_boot())
        write_binary(fd, 0x30a6c, get_lib()[:0x8df0])
    else:
        write_binary(fd, 0x30a20, get_lib()[:0x52])
        write_binary(fd, 0x2ea22, get_boot())
        write_binary(fd, 0x30a72, get_lib()[:0x8df0])
    write_binary(fd, 0x157120, get_unknown())

    write_binary(fd, 0, get_header())
    write_binary(fd, 0x40, get_rspboot())
    write_binary(fd, 0x1000, get_boot())
    write_binary(fd, 0x3050, get_lib())
    write_binary(fd, 0x39850, get_gamedata())
    write_binary(fd, 0x4e850, get_inflate())
    write_binary(fd, 0x4fc40, get_gamezips())
    write_binary(fd, 0x7d0a40, get_mpconfigs())
    write_binary(fd, 0x7d1c20, get_mpstrings('E'))
    write_binary(fd, 0x7d5320, get_mpstrings('J'))
    write_binary(fd, 0x7d8a20, get_mpstrings('P'))
    write_binary(fd, 0x7dc120, get_mpstrings('G'))
    write_binary(fd, 0x7df820, get_mpstrings('F'))
    write_binary(fd, 0x7e2f20, get_mpstrings('S'))
    write_binary(fd, 0x7e6620, get_mpstrings('I'))
    write_binary(fd, 0x7f2388, get_fonts())
    write_binary(fd, 0x80a250, get_sfxctl())
    write_binary(fd, 0x839dd0, get_sfxtbl())
    write_binary(fd, 0xcfbf30, get_seqctl())
    write_binary(fd, 0xd05f90, get_seqtbl())
    write_binary(fd, 0xe82000, get_midi())
    write_files(fd)
    write_binary(fd, 0x1d5ca00, get_filenames())
    write_binary(fd, 0x1d65f40, get_textures())
    fd.close()

def write_binary(fd, address, binary):
    fd.seek(address)
    fd.write(binary)

def get_header():
    binary = bytearray()
    binary.extend(b'\x80\x37\x12\x40') # Identifier
    binary.extend(b'\x00\x00\x00\x0f') # Clock rate
    binary.extend(b'\x80\x00\x10\x00') # Program counter
    binary.extend(b'\x00\x00\x14\x49') # Release address
    binary.extend(b'\x00\x00\x00\x00') # CRC 1
    binary.extend(b'\x00\x00\x00\x00') # CRC 2
    binary.extend(b'\x00\x00\x00\x00')
    binary.extend(b'\x00\x00\x00\x00')
    binary.extend(b'Perfect Dark        ')
    binary.extend(b'\x00\x00\x00\x00')
    binary.extend(b'\x00\x00\x00')
    binary.extend(b'NPDE')
    binary.extend(b'\x01' if os.environ['ROMID'] in ['ntsc-beta', 'ntsc-final'] else b'\x00')
    return binary

def get_rspboot():
    return getfilecontents(edir() + '/ucode/rspboot.bin')

def get_boot():
    return getfilecontents(bdir() + '/ucode/boot.bin')

def get_lib():
    return zip(bdir() + '/ucode/lib.bin')

def get_gamedata():
    return zip(bdir() + '/ucode/gamedata.bin')

def get_inflate():
    return getfilecontents(bdir() + '/ucode/inflate.bin')

def get_gamezips():
    return getfilecontents(bdir() + '/ucode/gamezips.bin')

def get_unknown():
    return getfrombaserom(0x157120, 0x69b268)

def get_mpconfigs():
    return getfilecontents(bdir() + '/ucode/mpconfigs.bin')

def get_mpstrings(lang):
    return getfilecontents(bdir() + '/ucode/mpstrings%s.bin' % lang)

def get_fonts():
    return getfrombaserom(0x7f2388, 0x17ec8)

def get_sfxctl():
    return getfilecontents(edir() + '/audio/sfx.ctl')

def get_sfxtbl():
    return getfilecontents(edir() + '/audio/sfx.tbl')

def get_seqctl():
    return getfilecontents(edir() + '/audio/music.ctl')

def get_seqtbl():
    return getfilecontents(edir() + '/audio/music.tbl')

def get_midi():
    return getfilecontents(edir() + '/audio/sequences.bin')

def write_files(fd):
    start = getlinkervariable('_filesSegmentRomStart')
    end = getlinkervariable('_filesSegmentRomEnd')

    write_binary(fd, 0xed83a0, getfromldbin(start, end - start))

def get_filenames():
    return getfilecontents(bdir() + '/ucode/filenames.bin')

def get_textures():
    return getfrombaserom(0x01d65f40, 0x29a0c0)

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

def getfrombaserom(offset, len):
    fd = open('pd.%s.z64' % os.environ['ROMID'], 'rb')
    fd.seek(offset)
    binary = fd.read(len)
    fd.close()
    return binary

def getfromldbin(offset, len):
    fd = open(bdir() + '/pd.bin', 'rb')
    fd.seek(offset)
    binary = fd.read(len)
    fd.close()
    return binary

def zip(filename):
    return subprocess.check_output(['tools/rarezip', filename])

def bdir():
    return 'build/%s' % os.environ['ROMID']

def edir():
    return 'extracted/%s' % os.environ['ROMID']

def getlinkervariable(varname):
    if 'TOOLCHAIN' in os.environ:
        cmd = '%s-objdump' % os.environ['TOOLCHAIN']
    else:
        cmd = 'mips64-elf-objdump'

    objdump = subprocess.check_output([cmd, bdir() + '/pd.elf', '-t']).decode('utf-8')

    matches = re.findall(r'^([0-9a-f]+) .*? %s$' % varname, objdump, re.MULTILINE)
    return int(matches[0], 16)

main()