from rich.console import Console
from rich.markdown import Markdown
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/audio-analyser"),
    "VideoAnalyser": grpc.insecure_channel("[::]:50051/video-analyser"),
    "TapeIrregularityClassifier": grpc.insecure_channel("[::]:50051/tape-irregularity-classifier"),
    "TapeAudioRestoration": grpc.insecure_channel("[::]:50051/tape-audio-restoration"),
    "Packager": grpc.insecure_channel("[::]:50051/packager"),
}

def run(console: Console):
    with grpc.insecure_channel("[::]:50051") as channel:
        audio_analyser = arp_pb2_grpc.AIMStub(channel)

        request = arp_pb2.InfoRequest()
        response = audio_analyser.getInfo(request)
        console.print("[bold]{}[/], v{}".format(response.title, response.version))
        console.print(Markdown(response.description))

        request = arp_pb2.JobRequest(
            working_dir="../data",
            files_name="BERIO100",
            index=2,
        )
        with console.status("[bold]Computing...", 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)

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


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