packager.py

#!./venv/bin/activate

"""
MPAI CAE-ARP Packager.

Implements MPAI CAE-ARP Packager Technical Specification, providing:
- Access Copy Files:Images
  1. Restored Audio Files.
  2. Editing List.
  3. Set of Irregularity Images in a .zip file.
  4. Irregularity File.
- Preservation Master Files:
  1. Preservation Audio File.
  2. Preservation Audio-Visual File where the audio has been replaced with the Audio of the Preservation Audio File
     fully synchronised with the video.
  3. Set of Irregularity Images in a .zip file.
  4. Irregularity File.
"""

import json
import os
import shutil
import sys
from argparse import ArgumentParser, RawTextHelpFormatter
from moviepy.editor import VideoFileClip, AudioFileClip

from mpai_cae_arp.files import File, FileType
from mpai_cae_arp.io import Color, Style, pprint

__author__ = "Nadir Dalla Pozza"
__copyright__ = "Copyright 2022, Audio Innova S.r.l."
__credits__ = ["Niccolò Pretto", "Nadir Dalla Pozza", "Sergio Canazza"]
__license__ = "GPL v3.0"
__version__ = "1.0.1"
__maintainer__ = "Nadir Dalla Pozza"
__email__ = "nadir.dallapozza@unipd.it"
__status__ = "Production"


def get_arguments() -> tuple[str, str]:
    """
    Method to obtain arguments from config/args.yaml file or command line.
    Default config/args.yaml, ignored if a command line argument is passed.
    :return: tuple consisting of two strings:
             1) the working path;
             2) the name of the Preservation files, which is key element to retrieve necessary files.
    """
    if len(sys.argv) > 1:
        # Read from command line
        parser = ArgumentParser(
            prog="python3 packager.py",
            formatter_class=RawTextHelpFormatter,
            description="A tool that implements MPAI CAE-ARP Packager Technical Specification.\n"
                        "By default, the configuration parameters are loaded from ./config/args.yaml file,\n"
                        "but, alternately, you can pass command line arguments to replace them."
        )
        parser.add_argument(
            "-w",
            "--working-path",
            help="Specify the Working Path, where all input files are stored",
            required=True
        )
        parser.add_argument(
            "-f",
            "--files-name",
            help="Specify the name of the Preservation files (without extension)",
            required=True
        )
        args = parser.parse_args()
        working_path = args.working_path
        files_name = args.files_name
    else:
        # Read from configuration file
        config = object
        try:
            config = File('./config/args.yaml', FileType.YAML).get_content()
            if 'WORKING_PATH' not in config:
                pprint('WORKING_PATH key not found in config/args.yaml!', color=Color.RED)
                quit(os.EX_CONFIG)
            if 'FILES_NAME' not in config:
                pprint('FILES_NAME key not found in config/args.yaml!', color=Color.RED)
                quit(os.EX_CONFIG)
        except FileNotFoundError:
            pprint('config/args.yaml file not found!', color=Color.RED)
            quit(os.EX_NOINPUT)
        working_path = config['WORKING_PATH']
        files_name = config['FILES_NAME']
    return working_path, files_name


def check_input(working_path: str, files_name: str) -> str:
    """
    Method to check that passed arguments are correct and that the environment is conformant to the standard.
    :param working_path: str representing the path where all files resulting from previous AIMs are stored.
    :param files_name: str representing the Preservation files name, to identify the input directory.
    :return: str representing the path where the files to be processed during the current execution are stored.
    """
    if not os.path.exists(working_path):
        pprint('The specified WORKING_PATH is non-existent!', color=Color.RED)
        quit(os.EX_CONFIG)
    # Check for temp directory existence
    temp_path = os.path.join(working_path, 'temp')
    if not os.path.exists(temp_path):
        pprint('WORKING_PATH structure is not conformant!', color=Color.RED)
        quit(os.EX_NOINPUT)
    # Check for input directory existence
    temp_path = os.path.join(temp_path, files_name)
    if not os.path.exists(temp_path):
        pprint('The specified FILES_NAME has no corresponding files!', color=Color.RED)
        quit(os.EX_NOINPUT)
    return temp_path


def make_dir(dir_path: str) -> bool:
    """
    Decide if the given directory has to be created or should be overridden.
    :param dir_path: str representing the path of the directory.
    :return: bool specifying if the files within the directory should be created.
    """
    if not os.path.exists(dir_path):
        # Create directory
        os.mkdir(dir_path)
        print("Directory '% s' created" % dir_path)
        return True
    else:
        pprint(f"Directory '{dir_path}' already exists!", color=Color.YELLOW)
        overwrite = input('Do you want to overwrite it? [y/n]: ')
        if overwrite.casefold() == 'y':
            # Overwrite directory
            shutil.rmtree(dir_path)
            os.mkdir(dir_path)
            print('%s overwritten' % dir_path)
            return True
        elif overwrite.casefold() != 'n':
            pprint('Unknown command, exiting', color=Color.RED)
            quit(os.EX_USAGE)
    return False


def main():
    """
    Main execution method.
    :return: exit codes corresponding to the execution status
    """
    pprint('\nWelcome to ARP Packager!', styles=[Style.BOLD])
    print(f"You are using Python version: {sys.version}")

    # Get the input from config/args.yaml or command line
    working_path, files_name = get_arguments()
    # Check if input is correct
    temp_path = check_input(working_path, files_name)

    # Access Copy Files
    pprint('\nCreation of Access Copy Files...', styles=[Style.BOLD])

    # By default, AccessCopyFiles directory is created in the output path...
    acf_dir = 'AccessCopyFiles/' + files_name + '/'
    acf_path = os.path.join(working_path, acf_dir)
    # ...however, if it already exists, the user can decide if it has to be overridden
    make_acf = make_dir(acf_path)

    if make_acf:
        # Copy RestoredAudioFiles
        raf_path = os.path.join(temp_path, 'RestoredAudioFiles')
        if not os.path.exists(raf_path):
            pprint(f"Restored Audio Files directory '{raf_path}' not found!", color=Color.RED)
            quit(os.EX_NOINPUT)
        restored_audio_files = os.listdir(raf_path)
        if len(restored_audio_files) == 1:
            pprint(f"Restored Audio Files directory '{raf_path}' is empty!", color=Color.YELLOW)
        shutil.copytree(raf_path, os.path.join(acf_path, 'RestoredAudioFiles'))
        print("Restored Audio Files copied")

        # Copy Editing List
        el_path = os.path.join(temp_path, 'EditingList.json')
        try:
            shutil.copy2(el_path, acf_path)
        except FileNotFoundError:
            pprint(f"Editing List file '{el_path}' not found!", color=Color.RED)
            quit(os.EX_NOINPUT)
        print("Editing List copied")

        # Create Irregularity Images archive
        ii_path = os.path.join(temp_path, 'IrregularityImages')
        if not os.path.exists(ii_path):
            pprint(f"Irregularity Images directory '{ii_path}' not found!", color=Color.RED)
            quit(os.EX_NOINPUT)
        irregularity_images = os.listdir(ii_path)
        if len(irregularity_images) == 1:
            pprint(f"Irregularity Images directory '{ii_path}' is empty!", color=Color.YELLOW)
        shutil.make_archive(acf_path + 'IrregularityImages', 'zip', temp_path, 'IrregularityImages')
        print("Irregularity Images archive created")

        # Copy Irregularity File
        if_path = os.path.join(temp_path, 'TapeIrregularityClassifier_IrregularityFileOutput2.json')
        try:
            shutil.copy2(if_path, os.path.join(acf_path, 'IrregularityFile.json'))
        except FileNotFoundError:
            pprint(f"Irregularity File file '{if_path}' not found!", color=Color.RED)
            quit(os.EX_NOINPUT)
        print("Irregularity File copied")

        # End Access Copy Files
        pprint('Success!', color=Color.GREEN, styles=[Style.BOLD])

    # Preservation Master Files
    pprint('\nCreation of Preservation Master Files...', styles=[Style.BOLD])

    # By default, PreservationMasterFiles directory is created in the output path...
    pmf_dir = 'PreservationMasterFiles/' + files_name + '/'
    pmf_path = os.path.join(working_path, pmf_dir)
    # ...however, if it already exists, it is possible to skip its creation
    make_pmf = make_dir(pmf_path)

    if make_pmf:

        # Copy Preservation Audio File
        audio_file = files_name + '.wav'
        paf_path = os.path.join(working_path, 'PreservationAudioFile/', audio_file)
        try:
            shutil.copy2(paf_path, pmf_path + 'PreservationAudioFile.wav')
        except FileNotFoundError:
            pprint(f"Preservation Audio File file '{paf_path}' not found!", color=Color.RED)
            quit(os.EX_NOINPUT)
        print("Preservation Audio File copied")

        # Create Preservation Audio-Visual File with substituted audio
        video_file = files_name + '.mov'
        pvf_path = os.path.join(working_path, 'PreservationAudioVisualFile/', video_file)
        try:
            audio = AudioFileClip(paf_path)
            video = VideoFileClip(pvf_path)
            # Open Irregularity File to get offset
            irregularity_file_json = open(
                os.path.join(temp_path, 'TapeIrregularityClassifier_IrregularityFileOutput2.json')
            )
            irregularity_file = json.load(irregularity_file_json)
            offset = irregularity_file['Offset']/1000
            if offset > 0:
                audio = audio.subclip(t_start=offset)
            else:
                video = video.subclip(t_start=offset)
            video = video.set_audio(audio)
            video.write_videofile(pmf_path + 'PreservationAudioVisualFile.mov', bitrate='3000k', codec='mpeg4')
            print("Preservation Audio-Visual File created")
        except OSError:
            pprint(f"Preservation Audio-Visual File file '{pvf_path}' not found!", color=Color.RED)
            quit(os.EX_NOINPUT)

        # Create Irregularity Images archive (already checked)
        shutil.make_archive(pmf_path + 'IrregularityImages', 'zip', temp_path, 'IrregularityImages')
        print("Irregularity Images archive created")

        # Copy Irregularity File (already checked)
        shutil.copy2(
            os.path.join(temp_path, 'TapeIrregularityClassifier_IrregularityFileOutput2.json'),
            os.path.join(pmf_path, 'IrregularityFile.json')
        )
        print("Irregularity File copied")

        # End Preservation Master Files
        pprint('Success!\n', color=Color.GREEN, styles=[Style.BOLD])

    else:
        print("\nExit")


if __name__ == '__main__':
    main()