import sys
from struct import unpack
from pathlib import Path
import traceback

# Constants
FILE_IDENTIFIER_SIZE = 4
UNKNOWN_BYTES_SIZE = 16
HEADER_AMOUNT_SIZE = 4
HEADER_LENGTH = 16
UNUSED_BYTES_SIZE = 4
NAME_ADDRESS_SIZE = 4
DATA_ADDRESS_SIZE = 4
DATA_LENGTH_SIZE = 4
BYTE_SIZE = 8

class Header():
    def __init__(self):
        self.unusedBytes = []
        self.nameAddress = 0
        self.dataAddress = 0
        self.dataSize = 0
        
        self.childHeader = None
        self.name = ""
        self.data = bytearray()
    
    def getHeaderInfoFromData(self, data, offset):
        assert(offset >= 0)
        assert(offset+HEADER_LENGTH <= len(data))
        pointer = 0
        self.unusedBytes = data[offset:offset+UNUSED_BYTES_SIZE]
        pointer += UNUSED_BYTES_SIZE
        self.nameAddress = readLittleEndian(data, offset+pointer)
        pointer += NAME_ADDRESS_SIZE
        self.dataAddress = readLittleEndian(data, offset+pointer)
        pointer += DATA_ADDRESS_SIZE
        self.dataSize = readLittleEndian(data, offset+pointer)
    
    def getName(self, data):
        assert(self.nameAddress >= 0)
        self.name = ""
        pointer = self.nameAddress
        while (data[pointer] != 0x00):
            self.name = self.name + chr(data[pointer])
            pointer += 1
    
    def getData(self, data):
        if self.isFolder():
            return
        assert(self.dataAddress >= 0)
        assert(self.dataSize >= 0)
        self.data.clear()
        for pointer in range(self.dataAddress, self.dataAddress+self.dataSize):
            self.data.append(data[pointer])
    
    def isFolder(self):
        if (self.dataSize == 0):
            return True
        return False

def readLittleEndian(data, offset):
    assert(offset >= 0)
    assert(offset+4 <= len(data))
    value = 0
    swapped = unpack('<L', data[offset:offset+4])
    value = swapped[0]
    return value
    
def getAllHeaders(data, offset):
    headers = []
    pointer = offset
    headerAmount = readLittleEndian(filedata, pointer)
    pointer += HEADER_AMOUNT_SIZE
    for headerNum in range(0, headerAmount):
        header = Header()
        header.getHeaderInfoFromData(filedata, pointer)
        pointer += HEADER_LENGTH
        header.getName(data)
        header.getData(data)
        if header.isFolder():
            header.childHeader = getAllHeaders(data, header.dataAddress)
        headers.append(header)
        print('.', end='', flush=True)
    return headers

def writeFilesFromHeaders(headers, path=Path('.')):
    filepath = path
    assert(filepath.exists())
    assert(filepath.is_dir())
    for header in headers:
        if header.name is not None:
            p = filepath / header.name
            if header.isFolder():
                if not p.exists():
                    p.mkdir()
                if p.exists():
                    writeFilesFromHeaders(header.childHeader, p)
            else:
                if header.data is not None:
                    pass
                    with p.open(mode='wb') as f:
                        f.write(header.data)
        print('.', end='', flush=True)

if __name__ == "__main__":
    f = None
    try:
        print("Opening file")
    
        f = open('umires.arc', 'rb')
        filedata = f.read()
        f.close()
        
        print("Validating file type")
        
        pointer = 0
        if ((filedata[pointer] != 0x53) or (filedata[pointer+1] != 0x41) 
           or (filedata[pointer+2] != 0x43) or (filedata[pointer+3] != 0x48)):
            raise TypeError("The file type is not a valid .arc file for the game")
        pointer += FILE_IDENTIFIER_SIZE
        
        # Unknown pack of 16 bytes after file type
        pointer += UNKNOWN_BYTES_SIZE
        
        print("Checking file headers", end='', flush=True)
        headers = getAllHeaders(filedata, pointer)
        print("")
        
        print("Extracting data from file", end='', flush=True)
        writeFilesFromHeaders(headers)
        print("")
        
        print("Files have been extracted. Have fun!")
    except Exception as inst:
        if f:
            f.close()
        exc_type, exc_value, exc_traceback = sys.exc_info()
        traceback.print_exception(exc_type, exc_value, exc_traceback, limit=8, file=sys.stdout)
