#! python3

import argparse
import io

from struct import unpack

BUFFER_SIZE = io.DEFAULT_BUFFER_SIZE


def main():
    """Entry point"""
    parser = argparse.ArgumentParser(description='NB0 Archive Unpacker')
    parser.add_argument('file', help='File to unpack')
    parser.add_argument('--overwrite', help='Overwrite extracted files',
                        dest='overwrite', action='store_true')

    args = parser.parse_args()

    if args.overwrite:
        output_file_mode = 'wb'
    else:
        output_file_mode = 'xb'

    with open(args.file, mode='rb') as archive:
        entries = {}

        data = archive.read(4)
        file_count = unpack('<i', data)[0]

        data_offset = 4 + file_count * 64

        print(f"I: Files in archive: {file_count}")

        for _ in range(file_count):
            data = archive.read(64)
            (offset, size, unknown1, unknown2, name) = unpack("<4I48s", data)
            name = name.decode('ascii').rstrip('\0')
            print(f"V: {name}: {size} bytes @ {offset} ({unknown1}, {unknown2})")

            if name not in entries:
                entries[name] = {'offset': offset, 'size': size}
            else:
                print(f'W: "{name}" already exists in archive.')

        for filename, info in entries.items():
            print(f"I: Extracting {filename} ({info['size']} bytes)...")
            with open(filename, mode=output_file_mode) as output:
                archive.seek(data_offset + info['offset'], 0)
                remaining_bytes = info['size']
                while remaining_bytes > 0:
                    chunk_size = min(BUFFER_SIZE, remaining_bytes)
                    output.write(archive.read(chunk_size))
                    remaining_bytes -= chunk_size
        print("I: Done")

if __name__ == "__main__":
    main()
