import argparse
import asyncio
from pathlib import Path

from riva.client.asr import AudioChunkFileIterator
from riva.client.realtime import RealtimeClientASR


NIM_SERVER = "127.0.0.1:9000"
INPUT_FILE = "/home/ubuntu/voice-assistant/audios/chafik_normalize.wav"
CHUNK_FRAMES = 1600


def build_args() -> argparse.Namespace:
    return argparse.Namespace(
        input_file=INPUT_FILE,
        mic=False,
        list_devices=False,
        duration=None,
        input_device=None,

        sample_rate_hz=16000,
        num_channels=1,
        file_streaming_chunk=CHUNK_FRAMES,

        output_text=None,
        prompt="",

        server=NIM_SERVER,
        ssl_root_cert=None,
        ssl_client_cert=None,
        ssl_client_key=None,
        use_ssl=False,
        metadata=None,
        options=None,
        max_message_length=67108864,

        word_time_offsets=False,
        max_alternatives=1,
        profanity_filter=False,
        automatic_punctuation=True,
        no_verbatim_transcripts=False,

        language_code="pt-BR",
        model_name="",

        boosted_lm_words=None,
        boosted_lm_score=4.0,

        speaker_diarization=False,
        diarization_max_speakers=3,

        start_history=-1,
        start_threshold=-1.0,
        stop_history=-1,
        stop_threshold=-1.0,
        stop_history_eou=-1,
        stop_threshold_eou=-1.0,

        custom_configuration="",
        endpoint="/v1/realtime",
        query_params="intent=transcription",
        mint_api_key=None,
    )


async def main() -> None:
    input_path = Path(INPUT_FILE)

    if not input_path.exists():
        raise FileNotFoundError(
            f"Audio file not found: {input_path}"
        )

    args = build_args()

    print(f"NIM server : {args.server}")
    print(f"Audio      : {input_path}")
    print(f"Language   : {args.language_code}")
    print(f"Chunk      : {CHUNK_FRAMES} frames")

    client = RealtimeClientASR(args=args)

    audio_iterator = AudioChunkFileIterator(
        input_file=str(input_path),
        chunk_n_frames=CHUNK_FRAMES,
        delay_callback=None,
    )

    try:
        print("Connecting to NIM...")
        await client.connect()
        print("Connected.")

        send_task = asyncio.create_task(
            client.send_audio_chunks(audio_iterator)
        )

        receive_task = asyncio.create_task(
            client.receive_responses()
        )

        await asyncio.gather(
            send_task,
            receive_task,
        )

    finally:
        await client.disconnect()
        print("Disconnected.")


if __name__ == "__main__":
    asyncio.run(main())