"""5.8 Transcribing while you speak: stream your microphone (or an audio file) to the AI Builder Space real-time speech endpoint and print words as they arrive.

Usage:
    export AI_BUILDER_TOKEN=sk_...
    python live.py                          # Use default microphone; press Enter when done to get finalized text
    python live.py --terms names.txt        # Pass a vocabulary list at the ASR stage (one term per line)
    python live.py --file my.wav            # Stream a 24 kHz mono 16-bit WAV at real-time speaking speed without a microphone
    python live.py --device 3               # Specify microphone device index (python -c "import sounddevice; print(sounddevice.query_devices())")
Microphone capture requires sounddevice and numpy (already included in requirements.txt).
"""
import argparse
import asyncio
import json
import os
import sys
import time
import wave
from pathlib import Path

import httpx
import websockets

BASE_URL = os.getenv('AI_BUILDER_SPACE_BASE_URL', 'https://space.ai-builders.com/backend/v1')
WS_BASE = BASE_URL.replace('https://', 'wss://').replace('http://', 'ws://').removesuffix('/v1')
RATE, CHUNK_SEC = 24000, 0.1


async def file_source(path, queue):
    with wave.open(path) as w:
        if (w.getframerate(), w.getnchannels(), w.getsampwidth()) != (RATE, 1, 2):
            raise SystemExit('Requires a 24 kHz mono 16-bit WAV; see the ffmpeg command in README.md')
        pcm = w.readframes(w.getnframes())
    step = int(RATE * CHUNK_SEC) * 2
    for i in range(0, len(pcm), step):
        await queue.put(pcm[i:i + step])
        await asyncio.sleep(CHUNK_SEC)
    await queue.put(None)                                    # Finished streaming file (end of speech)


async def mic_source(device, queue):
    import sounddevice as sd
    loop = asyncio.get_running_loop()
    stop = asyncio.Event()

    def callback(indata, frames, t, status):
        loop.call_soon_threadsafe(queue.put_nowait, bytes(indata))

    with sd.RawInputStream(samplerate=RATE, channels=1, dtype='int16', blocksize=int(RATE * CHUNK_SEC),
                           device=device, callback=callback):
        print('Start speaking; press Enter when finished...\n')
        await loop.run_in_executor(None, sys.stdin.readline)
        stop.set()
    await queue.put(None)


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--file')
    ap.add_argument('--device', type=int)
    ap.add_argument('--terms', help='Vocabulary file, one term per line')
    ap.add_argument('--language', default='zh')
    ap.add_argument('--model', default='gpt-live-transcribe')
    args = ap.parse_args()
    token = os.getenv('AI_BUILDER_TOKEN') or sys.exit('Please set the AI_BUILDER_TOKEN environment variable first')
    body = {'model': args.model, 'language': args.language, 'vad': False}
    if args.terms:
        body['terms'] = [l.strip() for l in Path(args.terms).read_text(encoding='utf-8').splitlines() if l.strip()]
    r = httpx.post(f'{BASE_URL}/audio/realtime/sessions', json=body, headers={'Authorization': f'Bearer {token}'}, timeout=60)
    r.raise_for_status()

    queue = asyncio.Queue()
    async with websockets.connect(WS_BASE + r.json()['ws_url'], max_size=None) as ws:
        t0 = time.time()
        source = asyncio.create_task(file_source(args.file, queue) if args.file else mic_source(args.device, queue))
        final = asyncio.get_running_loop().create_future()

        async def receive():
            async for raw in ws:
                msg = json.loads(raw)
                if msg['type'] == 'transcript_delta':
                    print(msg['text'], end='', flush=True)          # Print incremental text while speaking
                elif msg['type'] == 'transcript_completed':
                    final.set_result(msg['text'])
                elif msg['type'] == 'error':
                    final.set_exception(RuntimeError(msg.get('message')))
        rx = asyncio.create_task(receive())

        while (chunk := await queue.get()) is not None:
            await ws.send(chunk)
        t_end = time.time()
        await ws.send(json.dumps({'type': 'commit'}))
        text = await asyncio.wait_for(final, 120)
        print(f'\n\nFinal text (received {time.time() - t_end:.2f} s after speech ended, {time.time() - t0:.1f} s total):\n{text}')
        await ws.send(json.dumps({'type': 'stop'}))
        rx.cancel()
        await source


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