
import struct
import argparse
import os


def pack_pseudo_archive(output_file, input_files):
    with open(output_file, 'wb') as archive:
        count = len(input_files)
        archive.write(struct.pack('<H', count))  # 2 bytes for file count

        # Reserve space for offsets
        archive.write(b'    ' * count)

        offsets = []
        for file in input_files:
            offsets.append(archive.tell())
            with open(file, 'rb') as f:
                archive.write(f.read())

        # Go back and write offsets
        archive.seek(2)
        for offset in offsets:
            archive.write(struct.pack('<I', offset))


def unpack_pseudo_archive(archive_file):
    with open(archive_file, 'rb') as archive:
        count = struct.unpack('<H', archive.read(2))[0]
        offsets = [struct.unpack('<I', archive.read(4))[0] for _ in range(count)]

        offsets.append(os.path.getsize(archive_file))  # Add end of file for last file size

        for i in range(count):
            archive.seek(offsets[i])
            data = archive.read(offsets[i+1] - offsets[i])
            output_name = f"{archive_file}_{i+1:02d}.dbe"
            with open(output_name, 'wb') as out:
                out.write(data)


def import_file_into_archive(import_file, archive_file):
    # Extract index from filename
    base = os.path.basename(import_file)
    parts = base.split('_')
    if len(parts) < 2 or not parts[1][:2].isdigit():
        raise ValueError("Invalid import filename format. Expected format: archive_XX.uakm")

    index = int(parts[1][:2]) - 1  # Convert to 0-based index

    with open(archive_file, 'rb') as archive:
        count = struct.unpack('<H', archive.read(2))[0]
        offsets = [struct.unpack('<I', archive.read(4))[0] for _ in range(count)]
        offsets.append(os.path.getsize(archive_file))

        file_data = []
        for i in range(count):
            archive.seek(offsets[i])
            file_data.append(archive.read(offsets[i+1] - offsets[i]))

    # Replace the specified file
    with open(import_file, 'rb') as f:
        file_data[index] = f.read()

    # Repack the archive
    with open(archive_file, 'wb') as archive:
        archive.write(struct.pack('<H', count))
        archive.write(b'    ' * count)

        new_offsets = []
        for data in file_data:
            new_offsets.append(archive.tell())
            archive.write(data)

        archive.seek(2)
        for offset in new_offsets:
            archive.write(struct.pack('<I', offset))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Pseudo-archive utility')
    subparsers = parser.add_subparsers(dest='command')

    pack_parser = subparsers.add_parser('pack')
    pack_parser.add_argument('output', help='Output archive file')
    pack_parser.add_argument('inputs', nargs='+', help='Input files to pack')

    unpack_parser = subparsers.add_parser('unpack')
    unpack_parser.add_argument('archive', help='Archive file to unpack')

    import_parser = subparsers.add_parser('import')
    import_parser.add_argument('file', help='File to import into archive')
    import_parser.add_argument('archive', help='Target archive file')

    args = parser.parse_args()

    if args.command == 'pack':
        pack_pseudo_archive(args.output, args.inputs)
    elif args.command == 'unpack':
        unpack_pseudo_archive(args.archive)
    elif args.command == 'import':
        import_file_into_archive(args.file, args.archive)
    else:
        parser.print_help()
