import asyncio
import wave

from voice_assistant.stt.nvidia import NvidiaSTTAdapter


AUDIO_FILE = (
    "/home/ubuntu/voice-assistant/audios/chafik_normalize.wav"
)

CHUNK_FRAMES = 1600


async def main() -> None:
    stt = NvidiaSTTAdapter(
        server="127.0.0.1:9000",
        language_code="pt-BR",
        automatic_punctuation=True,
    )

    print("Starting STT session...")

    await stt.start_session()

    print("STT session ready.")

    async def send_audio() -> None:
        with wave.open(AUDIO_FILE, "rb") as wav:
            while True:
                audio = wav.readframes(CHUNK_FRAMES)

                if not audio:
                    break

                await stt.send_audio(audio)

    async def receive_events() -> None:
        async for event in stt.events():
            print(
                f"[{event.type.upper():7}] {event.text}"
            )

    sender = asyncio.create_task(send_audio())
    receiver = asyncio.create_task(receive_events())

    try:
        await sender
        print("Audio stream sent.")

        await stt.stop_session()

        await receiver

    finally:
        if not sender.done():
            sender.cancel()

        if not receiver.done():
            receiver.cancel()

        await asyncio.gather(
            sender,
            receiver,
            return_exceptions=True,
        )

    print("STT session finished.")


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