import os
import argparse
from typing import List
from reedsolo import RSCodec

def parse_arguments():
    parser = argparse.ArgumentParser(description='Reassemble file from multiple Reed-Solomon encoded pieces.')
    parser.add_argument('input_files', type=str, nargs='+', help='List of files containing the encoded pieces, in order.')
    return parser.parse_args()

def reassemble_file(input_files: List[str]):
    # Load the encoded pieces into memory
    encoded_pieces = []
    for input_file in input_files[:-1]:
        with open(input_file, 'rb') as f:
            encoded_piece_data = f.read()
            encoded_piece = bytearray(encoded_piece_data)
            print(f"Read {len(encoded_piece)} bytes from {input_file}")  # Print the size of the bytearray
            encoded_pieces.append(encoded_piece)

    # Decode the pieces using Reed-Solomon coding
    n = len(encoded_pieces)
    codec = RSCodec(n)
    decoded_pieces = codec.decode(encoded_pieces)

    # Concatenate the decoded pieces into the original file data
    file_data = b''.join(decoded_pieces)

    # Write the original file data to disk
    output_file_path = input_files[-1]
    with open(output_file_path, 'wb') as f:
        f.write(file_data)


if __name__ == '__main__':
    args = parse_arguments()
    reassemble_file(args.input_files)