cli.py

import argparse
import os
import sys
from rich.console import Console
import segment_finder as sf

from mpai_cae_arp.types.irregularity import IrregularityFile
from mpai_cae_arp.files import File, FileType
from mpai_cae_arp.io import prettify, Style


def get_args() -> tuple[str | None, str | None]:
    if len(sys.argv) > 1:
        parser = argparse.ArgumentParser(
            prog="audio-analyzer",
            formatter_class=argparse.RawDescriptionHelpFormatter,
            description=f"A tool that implements {prettify('MPAI CAE-ARP Audio Analyser', styles=[Style.BOLD])} Technical Specification.",
            epilog="For support, please contact Matteo Spanio <dev2@audioinnova.com>.\n"
                "This software is licensed under the GNU General Public License v3.0."
        )
        parser.add_argument("--working-directory", "-w", help="The path were the AIW will find and save the files")
        parser.add_argument("--files-name", "-f", help=f"The name of the files to be analyzed {prettify('without extension', styles=[Style.UNDERLINE])}")
        args = parser.parse_args()
        return args.working_directory, args.files_name
    else:
        return os.getenv("WORKING_DIRECTORY"), os.getenv("FILES_NAME")



def exit_with_error(error_message: str, console) -> None:
    console.print(f"[red bold]Error: {error_message}")
    quit(os.EX_USAGE)


def main() -> None:
    console = Console()
    console.print("[bold]Welcome to ARP Audio Analyser!")

    working_directory, files_name = get_args()
    if any(map(lambda x: x is None, [working_directory, files_name])):
        exit_with_error("{}\n{}".format(
            "Working directory or files name not specified!",
            "Try -h/--help to know more about Audio Analyser usage"), console)

    os.makedirs(os.path.join(working_directory, "temp", files_name))
    
    with console.status("[purple]Reading input files", spinner="dots"):
        audio_src = os.path.join(working_directory, "PreservationAudioFile", f"{files_name}.wav")
        video_src = os.path.join(working_directory, "PreservationAudioVisualFile", f"{files_name}.mov")

        audio_exists = os.path.exists(audio_src)
        video_exists = os.path.exists(video_src)

        match audio_exists, video_exists:
            case True, True:
                console.print("[green]Input files found!")
            case False, True:
                exit_with_error("Audio file not found! :loud_sound:", console)
            case True, False:
                exit_with_error("Video file not found! :vhs:", console)
            case False, False:
                exit_with_error("Input files not found! :t-rex:", console)

    # create irregularity file 1
    with console.status("[purple]Creating irregularity file 1", spinner="dots"):
        irreg1 = sf.create_irreg_file(audio_src, video_src)
        console.log(f"Found {len(irreg1.irregularities)} irregularities from Audio source")
        File(f"{working_directory}/temp/{files_name}/IrregularityFile1.json", FileType.JSON).write_content(irreg1.to_json())
        console.log("[geen]Irregularity file 1 created")

    # create irregularity file 2
    with console.status("[purple]Creating irregularity file 2", spinner="dots"):
        video_irreg_1 = {
            "Irregularities": [
                {
                    "TimeLabel": "00:03:00.000",
                    "Source": "v",
                    "IrregularityID": "09a0b0c0-d0e0-f000-0000-000000000000"
                },
                {
                    "TimeLabel": "00:03:01.000",
                    "Source": "v",
                    "IrregularityID": "09a0b0c0-d0e0-f000-0000-000000000001"
                }
            ],
            "Offset": 170
            }
        console.log("Video irregularity file 1 found")
        irreg2 = sf.merge_irreg_files(irreg1, IrregularityFile.from_json(video_irreg_1))
        console.log("[geen]Irregularity file 2 created")

    with console.status("[cyan]Extracting audio irregularities", spinner="dots"):
        irreg2= sf.extract_audio_irregularities(audio_src, irreg2, working_directory + "/temp/" + files_name)
        File(f"{working_directory}/temp/{files_name}/IrregularityFile2.json", FileType.JSON).write_content(irreg2.to_json())
        console.log("[green]Audio irregularities extracted")

    # classify audio irregularities
    with console.status("[cyan bold]Classifying audio irregularities", spinner="monkey"):
        sf.classify_audio_irregularities(working_directory)

    console.print("[green bold]Success! :tada:")
    quit(os.EX_OK)


if __name__ == "__main__":
    main()