from rich.console import Console
import os

import grpc
from mpai_cae_arp.network import arp_pb2
from mpai_cae_arp.network import arp_pb2_grpc

channels = {
    "AudioAnalyser": grpc.insecure_channel("[::]:50051"),
    "VideoAnalyser": grpc.insecure_channel("[::]:50052"),
    "TapeIrregularityClassifier": grpc.insecure_channel("[::]:50053"),
    "TapeAudioRestoration": grpc.insecure_channel("[::]:50051"),
    #"Packager": grpc.insecure_channel("[::]:50051/packager"),
}

def run(console: Console):

    audio_analyser = arp_pb2_grpc.AIMStub(channels["AudioAnalyser"])
    video_analyser = arp_pb2_grpc.AIMStub(channels["VideoAnalyser"])
    tape_irreg_classifier = arp_pb2_grpc.AIMStub(channels["TapeIrregularityClassifier"])
    tape_audio_restoration = arp_pb2_grpc.AIMStub(channels["TapeAudioRestoration"])

    request = arp_pb2.InfoRequest()
    # for analyser in [audio_analyser, video_analyser, tape_irreg_classifier]:
    #     response = analyser.getInfo(request)
    #     console.print("[bold]{}[/], v{}".format(response.title, response.version))

    request = arp_pb2.JobRequest(
        working_dir="../data",
        files_name="BERIO100",
        index=1,
    )

    # with console.status("[bold]Computing AudioAnalyser IrregularityFile 1...", spinner="bouncingBall"):
    #     for result in audio_analyser.work(request):
    #         if result.status == "error":
    #             console.print("[bold red]Error![/] :boom:")
    #             console.print(f"[italic red]{result.message}")
    #             for channel in channels.values():
    #                 channel.close()
    #             exit(os.EX_SOFTWARE)
    #         console.print(result.message)

    # request.files_name = "BERIO100.mov"
    # with console.status("[bold]Computing VideoAnalyser IrregularityFiles...", spinner="bouncingBall"):
    #     for result in video_analyser.work(request):
    #         if result.status == "error":
    #             console.print("[bold red]Error![/] :boom:")
    #             console.print(f"[italic red]{result.message}")
    #             for channel in channels.values():
    #                 channel.close()
    #             exit(os.EX_SOFTWARE)
    #         console.print(result.message)

    # request.index = 2
    # request.files_name = "BERIO100"
    # with console.status("[bold]Computing AudioAnalyser IrregularityFile 2...", spinner="bouncingBall"):
    #     for result in audio_analyser.work(request):
    #         if result.status == "error":
    #             console.print("[bold red]Error![/] :boom:")
    #             console.print(f"[italic red]{result.message}")
    #             for channel in channels.values():
    #                 channel.close()
    #             exit(os.EX_SOFTWARE)
    #         console.print(result.message)

    # with console.status("[bold]Computing TapeIrregularityClassifier...", spinner="bouncingBall"):
    #     for result in tape_irreg_classifier.work(request):
    #         if result.status == "error":
    #             console.print("[bold red]Error![/] :boom:")
    #             console.print(f"[italic red]{result.message}")
    #             for channel in channels.values():
    #                 channel.close()
    #             exit(os.EX_SOFTWARE)
    #         console.print(result.message)

    with console.status("[bold]Computing TapeAudioRestoration...", spinner="bouncingBall"):
        for result in tape_audio_restoration.work(request):
            if result.status == "error":
                console.print("[bold red]Error![/] :boom:")
                console.print(f"[italic red]{result.message}")
                for channel in channels.values():
                    channel.close()
                exit(os.EX_SOFTWARE)
            console.print(result.message)

    channels["AudioAnalyser"].close()
    channels["VideoAnalyser"].close()
    channels["TapeIrregularityClassifier"].close()
    channels["TapeAudioRestoration"].close()

    console.print("[bold green]Success![/] :tada:")


if __name__ == '__main__':
    console = Console()
    run(console)