import os
import argparse
from math import ceil
from pydantic import BaseModel
from typing import List
from reedsolo import RSCodec

class Args(BaseModel):
    file_path: str
    n: int

def parse_arguments():
    parser = argparse.ArgumentParser(description='Split file into multiple pieces using Reed-Solomon coding.')
    parser.add_argument('file_path', type=str, help='Path to the input file.')
    parser.add_argument('n', type=int, help='Number of pieces to split the file into.')
    return Args(**vars(parser.parse_args()))

def split_file_into_pieces(file_path: str, n: int):
    # Load the file into memory
    with open(file_path, 'rb') as f:
        file_data = f.read()

    # Calculate the size of each piece
    piece_size = ceil(len(file_data) / n)

    # Create the folders for the pieces
    folder_names = [f'piece_{i+1}' for i in range(n)]
    for folder_name in folder_names:
        os.makedirs(folder_name, exist_ok=True)

    # Encode each piece using Reed-Solomon coding and write it to the corresponding folder
    codec = RSCodec(n)
    for i, folder_name in enumerate(folder_names):
        piece_data = file_data[i*piece_size:(i+1)*piece_size]
        encoded_data = codec.encode(piece_data)
        with open(os.path.join(folder_name, f'{i+1}.dat'), 'wb') as f:
            f.write(encoded_data)

def reassemble_file(folder_names: List[str], output_file_path: str):
    # Load the encoded pieces into memory
    encoded_pieces = []
    for i, folder_name in enumerate(folder_names):
        with open(os.path.join(folder_name, f'{i+1}.dat'), 'rb') as f:
            encoded_piece_data = f.read()
            encoded_pieces.append(encoded_piece_data)

    # Decode the pieces using Reed-Solomon coding
    codec = RSCodec(len(encoded_pieces))
    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
    with open(output_file_path, 'wb') as f:
        f.write(file_data)

if __name__ == '__main__':
    args = parse_arguments()
    split_file_into_pieces(args.file_path, args.n)
    folder_names = [f'piece_{i+1}' for i in range(args.n)]
    reassemble_file(folder_names, 'reconstructed_file.png')