import os
from os import path
import sys
from argparse import ArgumentParser, RawTextHelpFormatter
from typing import NoReturn
from rich.console import Console

from mpai_cae_arp.types.irregularity import IrregularityType
from mpai_cae_arp.files import File, FileType

from tape_irregularity_classifier.lib import (
    CLASSES,
    collect_irregularity_images,
    get_irregularity_file,
    load_model,
    MODEL_PATH,
    predict,
    prepare_images_for_prediction,
    PROB,
    verify_path,
)

AUDIO_IRREG_FILE = 'AudioAnalyser_IrregularityFileOutput2.json'
VIDEO_IRREG_FILE = 'VideoAnalyser_IrregularityFileOutput2.json'
TAPE_IRREG_FILE_1 = 'TapeIrregularityClassifier_IrregularityFileOutput1.json'
TAPE_IRREG_FILE_2 = 'TapeIrregularityClassifier_IrregularityFileOutput2.json'


def get_args() -> tuple[str, str]:
    """
    Method to obtain arguments from environment variables or command line.
    Default environment, ignored if a command line argument is passed.

    Raises
    ------
    ValueError
        If no arguments are passed and environment variables are not set.

    Returns
    -------
    tuple[str, str]
        The working path and the name of the Preservation files, which is key element to retrieve necessary files.
    """
    if len(sys.argv) > 1:
        parser = ArgumentParser(
            prog='tapeIrregularityClassifier',
            formatter_class=RawTextHelpFormatter,
            description='A tool that implements MPAI CAE-ARP Tape Irregularity Classifier 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:
        working_path = os.getenv('WORKING_DIRECTORY')
        files_name = os.getenv('FILES_NAME')
        if any([working_path is None, files_name is None]):
            raise ValueError("Working path or files name not specified!")
    return working_path, files_name


def exit_with_error(error_message: str, error: int, console: Console) -> NoReturn:
    console.print(f"[red bold]Error: {error_message}")
    quit(error)


def main():
    """
    Execute the Tape Irregularity Classifier workflow:
    1. Get arguments from environment variables or command line.
    2. Verify that the environment is conformant to the standard.
    3. Read the irregularities from VideoAnalyser_IrregularityFileOutput2.json.
    4. Load the pretrained model.
    5. Collect the irregularity images to be classified.
    6. Classify the images.
    7. Write the results to TapeIrregularityClassifier_IrregularityFileOutput1.json and TapeIrregularityClassifier_IrregularityFileOutput2.json.

    Exceptions are handled and the program exits with the appropriate error code.
    """
    console = Console()

    console.print("[bold]Welcome to MPAI CAE-ARP TapeIrregularityClassifier!")
    console.print(f"You are using Python version: {sys.version}")

    try:
        working_directory, files_name = get_args()
        temp_path = verify_path(working_directory, files_name)
        irregularity_file = get_irregularity_file(path.join(temp_path, VIDEO_IRREG_FILE))
        audio_irregularity_file = get_irregularity_file(path.join(temp_path, AUDIO_IRREG_FILE))
        irregularity_file.offset = audio_irregularity_file.offset

        console.print(f"{len(irregularity_file.irregularities)} Irregularities listed in {VIDEO_IRREG_FILE}")

        model = load_model(MODEL_PATH)

        with console.status("Collecting irregularity images :camera:", spinner="dots"):
            irregularity_images = collect_irregularity_images(irregularity_file.irregularities)

        console.print(f"{len(irregularity_images)} images loaded.")
        images = prepare_images_for_prediction(irregularity_images)
        probabilities, label = predict(model, images)

        predicted_labels = []
        # The following array will correspond to the Irregularities that will be listed in the JSON
        output_irregularities = []

        for i in range(0, len(irregularity_file.irregularities)):
            # Retrieve current Irregularity probability
            irr_prob = probabilities[i][label[i]] * 100
            predicted_labels.append(CLASSES[label[i]] + ' - ' + str(irr_prob) + ' %')
            irr = irregularity_file.irregularities[i]
            # Consider valid Irregularities only those with probability > prob
            if probabilities[i][label[i]] * 100 > PROB:
                irr.irregularity_type = IrregularityType(CLASSES[label[i]])
                output_irregularities.append(irr)
            else:
                console.print(f"[cyan]Irregularity number {i+1} with probability {irr_prob:.2f} % discarded")
                console.print(f"[cyan]    ImageURI: {irr.image_URI}")
                console.print(f"[cyan]    IrregularityID: {irr.irregularity_ID}")
                # TODO: Delete corresponding image from temp_path
                # os.remove(irr['ImageURI'])

        console.print(f"Found {len(output_irregularities)} interesting Irregularities.")

        # TODO: IrregularityFileOutput1 shall present only Irregularities relevant to Tape Audio Restoration
        console.print(f"Saving {TAPE_IRREG_FILE_1} ...")
        File(path.join(temp_path, TAPE_IRREG_FILE_1), FileType.JSON)\
            .write_content(irregularity_file)
        console.print(f"Saving {TAPE_IRREG_FILE_2} ...")
        File(path.join(temp_path, TAPE_IRREG_FILE_2), FileType.JSON)\
            .write_content(irregularity_file)

    except FileNotFoundError as e:
        exit_with_error(e.strerror, os.EX_NOTFOUND, console)
    except ValueError as e:
        exit_with_error(
            "{}\n{}".format(
                "Working directory or files name not specified!",
                "Try -h/--help to know more about Tape Irregularity Classifier usage"),
            os.EX_USAGE, console)
    except Exception as e:
        exit_with_error(e, os.EX_SOFTWARE, console)

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


if __name__ == '__main__':
    main()
