"""5.8 边说边出字：把麦克风（或一段录音）实时推给 AI Builder Space 的实时语音接口，文字边出边打印。

用法：
    export AI_BUILDER_TOKEN=sk_...
    python live.py                          # 用默认麦克风；说完按回车，拿最终文字
    python live.py --terms names.txt        # 识别时带上名单（每行一个词）
    python live.py --file my.wav            # 不用麦克风，按真实语速推一段 24 kHz 单声道 16 位 WAV
    python live.py --device 3               # 指定麦克风编号（python -c "import sounddevice; print(sounddevice.query_devices())"）
用麦克风需要另外装 sounddevice 和 numpy（requirements 里已经有）。
"""
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('需要 24 kHz 单声道 16 位 WAV，见 README 里的 ffmpeg 命令')
        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)                                    # 录音放完，相当于说完了


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('开始说话，说完按回车……\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='名单文件，每行一个词')
    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('请先设置环境变量 AI_BUILDER_TOKEN')
    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)          # 边说边出字
                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\n最终文字（说完后 {time.time() - t_end:.2f} 秒拿到，全程 {time.time() - t0:.1f} 秒）：\n{text}')
        await ws.send(json.dumps({'type': 'stop'}))
        rx.cancel()
        await source


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