'''
Script to run tapeAudioRestoration.py on all files in a directory

Author: Matteo Spanio
Date: 2023-053-28

Usage:
    python script.py -i <input_directory> -o <output_directory>

the -o argument is optional, if not specified the script will use the WORKING_PATH variable in config.yaml
the -i argument is mandatory
the speed and standard values are taken from config.yaml

Description:
    Given a directory containing all the files to be processed, the script will create a folder structure in the output directory
    containing the original files in the PreservationAudioFile folder and the output of the tapeAudioRestoration.py script in the
    RestoredAudioFiles folder. The output folder will contain all the files with the same name as the original files.
'''

import os
import subprocess
import shutil
import yaml
import argparse


def create_folder_structure(destination, file_names, original_files):
    if os.path.exists(destination):
        shutil.rmtree(destination)
    os.makedirs(os.path.join(destination, "PreservationAudioFile"))
    for name in file_names:
        os.makedirs(os.path.join(destination, "temp", name))
    for file in original_files:
        shutil.copy(os.path.join(file), os.path.join(destination, "PreservationAudioFile"))

    
def get_files_name(input_dir):
    file_names = []
    for filename in os.listdir(input_dir):
        if filename.endswith(".wav"):
            file_names.append(filename.split(".")[0])
    return file_names


def read_config():
    with open("config.yaml", 'r') as stream:
        try:
            config = yaml.safe_load(stream)
        except yaml.YAMLError as exc:
            print(exc)
    return config


def main():
    parser = argparse.ArgumentParser(description="Script to run tapeAudioRestoration.py on all files in a directory")
    parser.add_argument("-i", "--input-directory", help="Directory containing all files to be processed", required=True)
    parser.add_argument("-o", "--output-directory", help="Directory to store all processed files")

    config = read_config()

    args = parser.parse_args()

    input_directory = os.path.abspath(args.input_directory)
    output_directory = args.output_directory or config["WORKING_PATH"]
    standard_w = config["STANDARD_W"]
    standard_r = config["STANDARD_R"]
    speed_w = str(config["SPEED_W"])
    speed_r = str(config["SPEED_R"])
    
    files_names = get_files_name(input_directory)

    create_folder_structure(output_directory, files_names, [os.path.join(input_directory, x) for x in os.listdir(input_directory)])

    for filename in files_names:
        subprocess.run(["python3", "./tapeAudioRestoration.py", "-w", output_directory, "-f", filename, "-ew", standard_w, "-sw", speed_w, "-sr", speed_r, "-er", standard_r])

    os.mkdir(os.path.join(output_directory, "RestoredAudioFiles"))
    for dir in os.listdir(os.path.join(output_directory, "temp")):
        shutil.move(os.path.join(output_directory, "temp", dir, "RestoredAudioFiles", "1.wav"), os.path.join(output_directory, "RestoredAudioFiles", dir + ".wav"))

    shutil.rmtree(os.path.join(output_directory, "temp"))

if __name__ == '__main__':

    main()