"""WebSocket transcription example. Requires Python 3.11+ and websockets 14–15."""
import asyncio
import base64
import json
import os

from websockets.asyncio.client import connect

URL = "wss://api.narilabs.com/v1/realtime?intent=transcription"


class InsufficientCredits(RuntimeError):
    """Stop this session; the caller must obtain credit before reconnecting."""


async def pcm_file(path):
    with open(path, "rb") as audio:
        while chunk := audio.read(3200):  # 100 ms at 16 kHz PCM16 mono
            yield chunk
            await asyncio.sleep(len(chunk) / 32000)


async def transcribe(audio_source, *, turn_detection=None, word_timestamps=False, on_event=print):
    """Source yields PCM bytes, or None to manually commit an utterance."""
    async with connect(
        URL,
        additional_headers={"Authorization": f"Bearer {os.environ['NARI_API_KEY']}"},
    ) as socket:
        async def receive():
            event = json.loads(await socket.recv())
            if event["type"] == "error":
                on_event(event)  # Preserve the event/request IDs for support.
                if event["error"]["code"] == "INSUFFICIENT_CREDITS":
                    raise InsufficientCredits("Add credit before starting another session")
                raise RuntimeError(event["error"])
            return event

        await socket.send(json.dumps({
            "type": "session.configure",
            "session": {
                "model": "qwen3-asr-fast",
                "turn_detection": turn_detection,
                "word_timestamps": word_timestamps,
            },
        }))
        assert (await receive())["type"] == "session.configured"

        async def send_audio():
            commit_number = 0
            async for chunk in audio_source:
                if chunk is None:
                    commit_number += 1
                    event = {"type": "input_audio_buffer.commit",
                             "event_id": f"commit_{commit_number}"}
                else:
                    if not chunk or len(chunk) % 2:
                        raise ValueError("Send nonempty chunks of whole PCM16 samples")
                    event = {"type": "input_audio_buffer.append",
                             "audio": base64.b64encode(chunk).decode("ascii")}
                await socket.send(json.dumps(event))
            # Flush any remaining audio, including when using VAD.
            await socket.send(json.dumps({
                "type": "input_audio_buffer.commit", "event_id": "end_of_input",
            }))

        async def read_results():
            pending = set()
            pending_words = set()
            end_acknowledged = False
            while True:
                event = await receive()
                kind = event["type"]
                if kind == "input_audio_buffer.committed":
                    pending.add(event["item_id"])
                elif kind == "transcript.completed":
                    pending.discard(event["item_id"])
                    if word_timestamps:
                        pending_words.add(event["item_id"])
                elif kind in ("transcript.words", "transcript.words.failed"):
                    pending_words.discard(event["item_id"])
                if kind in ("input_audio_buffer.committed", "input_audio_buffer.commit_empty"):
                    if event.get("client_event_id") == "end_of_input":
                        end_acknowledged = True
                on_event(event)
                if end_acknowledged and not pending and not pending_words:
                    return

        async with asyncio.TaskGroup() as tasks:
            tasks.create_task(send_audio())
            tasks.create_task(read_results())


if __name__ == "__main__":
    try:
        asyncio.run(transcribe(pcm_file("hello.pcm")))
    except* InsufficientCredits:
        print("Transcription stopped: add credit before reconnecting. Keep received final transcripts.")