"""5.8 实时识别：人名这类上下文，放进识别模型，还是放进后面的纠错模型？

对模块 2 的示例录音（每条录音有脚本原文作标准答案），用 AI Builder Space 的实时语音接口按真实语速推流，比较：
  plain    识别时不给名单
  terms    识别时把名单放进 terms
  correct  识别时不给名单，识别完再把转写和名单一起交给一个纠错模型（deepseek）
  both     识别时给 terms，识别完再交给纠错模型
  batch    （参考）说完以后把整段录音上传给整段识别接口，不给名单
每种做法数人名写对了几处，以及“说完到拿到最终文字”的延迟。

用法：
    export AI_BUILDER_TOKEN=sk_...
    python run.py --lang zh            # 默认：23 条含人名的录音，按真实语速推流，约 10 分钟
    python run.py --lang en
    python run.py --audio my_dir --memos my_memos.json   # 换成你自己的录音和标准答案，见 README
"""
import argparse
import asyncio
import io
import json
import os
import re
import statistics
import time
import urllib.request
import wave
import zipfile
from datetime import datetime
from pathlib import Path

import httpx
import websockets
from openai import AsyncOpenAI

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')
ZIP_URL = 'https://www.ai-builders.com/resources/aibuilder2027/m02/{lang}/ch02_memos_full38_{lang}.zip'
HERE = Path(__file__).resolve().parent
CHUNK = 4800          # 每次发 0.1 秒音频：24000 采样/秒 × 2 字节 × 0.1 秒

CORRECT_PROMPT = {
    'zh': '下面是一段语音备忘录的自动转写，可能有识别错误。请改正识别错误，只输出改好的全文，不要解释。\n\n'
          '这是我常提到的人和术语，转写里听起来像的地方请按这里的写法改：\n{names}\n\n转写：\n{text}',
    'en': 'Below is an automatic transcript of a voice memo. It may contain recognition errors. Fix them and output '
          'only the corrected full text, with no explanation.\n\nThese are people and terms I often mention; where the '
          'transcript has something that sounds like one of them, use this spelling:\n{names}\n\nTranscript:\n{text}',
}


def token():
    t = os.getenv('AI_BUILDER_TOKEN')
    if not t:
        raise SystemExit('请先设置环境变量 AI_BUILDER_TOKEN（你的 AI Builder Space key）')
    return t


def ensure_audio(lang, audio_dir):
    """第一次运行时下载模块 2 的 38 条示例录音（约 13 MB）。"""
    if audio_dir.exists() and any(audio_dir.glob('*.wav')):
        return
    print(f'下载示例录音到 {audio_dir.name}/ ……')
    data = urllib.request.urlopen(ZIP_URL.format(lang=lang)).read()
    zipfile.ZipFile(io.BytesIO(data)).extractall(audio_dir)


def read_pcm(path):
    with wave.open(str(path)) as w:
        if (w.getframerate(), w.getnchannels(), w.getsampwidth()) != (24000, 1, 2):
            raise SystemExit(f'{path.name} 不是 24 kHz 单声道 16 位 WAV，先用 ffmpeg 转换，见 README')
        return w.readframes(w.getnframes())


async def realtime(pcm, lang, model, terms):
    """推流一条录音。返回 (最终文字, 从开始推流到第一段文字的秒数, 从说完到最终文字的秒数)。"""
    body = {'model': model, 'language': lang, 'vad': False}
    if terms:
        body['terms'] = terms
    async with httpx.AsyncClient(timeout=60) as h:
        r = await h.post(f'{BASE_URL}/audio/realtime/sessions', json=body,
                         headers={'Authorization': f'Bearer {token()}'})
        r.raise_for_status()
    async with websockets.connect(WS_BASE + r.json()['ws_url'], max_size=None) as ws:
        t0 = time.time()
        first, done = None, asyncio.get_running_loop().create_future()

        async def receive():
            nonlocal first
            async for raw in ws:
                msg = json.loads(raw)
                if msg['type'] == 'transcript_delta' and first is None:
                    first = time.time() - t0
                elif msg['type'] == 'transcript_completed':
                    done.set_result((msg['text'], time.time()))
                    return
                elif msg['type'] == 'error':
                    done.set_exception(RuntimeError(msg.get('message')))
                    return

        rx = asyncio.create_task(receive())
        for i in range(0, len(pcm), CHUNK):
            await ws.send(pcm[i:i + CHUNK])
            await asyncio.sleep(0.1)              # 按真实语速发，模拟有人正在说话
        t_end = time.time()                       # 说完了
        await ws.send(json.dumps({'type': 'commit'}))
        text, t_final = await asyncio.wait_for(done, 120)
        await ws.send(json.dumps({'type': 'stop'}))
        rx.cancel()
        return text, first, t_final - t_end


async def batch(path, lang):
    """整段识别：说完以后把整个文件传上去。返回 (文字, 从上传到拿到文字的秒数)。"""
    t0 = time.time()
    async with httpx.AsyncClient(timeout=300) as h:
        r = await h.post(f'{BASE_URL}/audio/grok-transcription', headers={'Authorization': f'Bearer {token()}'},
                         data={'simple': 'true', 'language': lang},
                         files={'audio_file': (path.name, path.read_bytes())})
        r.raise_for_status()
    return r.json()['text'], time.time() - t0


async def correct(llm, text, names_md, lang, model):
    t0 = time.time()
    r = await llm.chat.completions.create(model=model, temperature=0, messages=[
        {'role': 'user', 'content': CORRECT_PROMPT[lang].format(names=names_md, text=text)}])
    return r.choices[0].message.content.strip(), time.time() - t0


def count(text, word, lang):
    """一个词在文本里出现几次。中文：忽略大小写和空格的子串计数；英文：忽略大小写的整词匹配。"""
    if lang == 'zh':
        norm = lambda s: s.lower().replace(' ', '')
        return norm(text).count(norm(word))
    body = r'[\s-]*'.join(re.escape(w) for w in word.split())
    return len(re.findall(rf'\b{body}\b', text, re.I))


def score(outputs, memos, vocab, lang):
    """某个词在脚本里出现 n 次，转写里出现 k 次，就算对 min(n, k) 处。"""
    total = hit = 0
    misses = {}
    for m in memos:
        for v in vocab:
            n = count(m['text'], v, lang)
            if n:
                k = min(n, count(outputs.get(m['id'], ''), v, lang))
                total, hit = total + n, hit + k
                if k < n:
                    misses[v] = misses.get(v, 0) + n - k
    return hit, total, misses


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--lang', choices=['zh', 'en'], default='zh')
    ap.add_argument('--model', default='gpt-live-transcribe', help='实时识别模型：gpt-live-transcribe 或 gpt-realtime')
    ap.add_argument('--llm', default='deepseek', help='纠错用的模型')
    ap.add_argument('--memos', help='标准答案文件（默认 memos_<lang>.json）')
    ap.add_argument('--audio', help='录音目录（默认 audio_<lang>/，没有就自动下载示例录音）')
    ap.add_argument('--only', help='只跑这些录音，逗号分隔，例如 W01,W02')
    ap.add_argument('--conditions', default='plain,terms,correct,both,batch')
    ap.add_argument('--workers', type=int, default=2, help='同时开几个实时会话')
    ap.add_argument('--out', default='results')
    args = ap.parse_args()

    d = json.loads(Path(args.memos or HERE / f'memos_{args.lang}.json').read_text(encoding='utf-8'))
    names = list(d['people'])
    terms = [t for t in d.get('terms', []) if t not in ('Apple', 'terminal')]
    vocab_list = names + terms                                  # 识别和纠错用的是同一份名单
    names_md = '\n'.join(f'- {k}：{v}' if args.lang == 'zh' else f'- {k}: {v}' for k, v in d['people'].items())
    names_md += ('\n\n常用术语：' + '、'.join(terms)) if args.lang == 'zh' else ('\n\nTerms: ' + ', '.join(terms))
    audio_dir = Path(args.audio) if args.audio else HERE / f'audio_{args.lang}'
    if not args.audio:
        ensure_audio(args.lang, audio_dir)
    memos = [m for m in d['memos'] if any(count(m['text'], n, args.lang) for n in names)]   # 只用提到人名的录音
    if args.only:
        memos = [m for m in memos if m['id'] in args.only.split(',')]
    conds = args.conditions.split(',')
    llm = AsyncOpenAI(base_url=BASE_URL, api_key=token(), timeout=300)
    sem = asyncio.Semaphore(args.workers)
    rows = {c: {} for c in conds}

    async def one(m):
        async with sem:
            pcm = read_pcm(audio_dir / f'{m["id"]}.wav')
            rec = {'seconds_of_audio': round(len(pcm) / 48000, 2)}
            if 'plain' in conds or 'correct' in conds:
                text, first, lat = await realtime(pcm, args.lang, args.model, None)
                rec['plain'] = dict(text=text, first_text_s=first, after_speech_s=lat)
            if 'terms' in conds or 'both' in conds:
                text, first, lat = await realtime(pcm, args.lang, args.model, vocab_list)
                rec['terms'] = dict(text=text, first_text_s=first, after_speech_s=lat)
        if 'correct' in conds:
            fixed, sec = await correct(llm, rec['plain']['text'], names_md, args.lang, args.llm)
            rec['correct'] = dict(text=fixed, first_text_s=rec['plain']['first_text_s'],
                                  after_speech_s=rec['plain']['after_speech_s'] + sec, llm_s=sec)
        if 'both' in conds:
            fixed, sec = await correct(llm, rec['terms']['text'], names_md, args.lang, args.llm)
            rec['both'] = dict(text=fixed, first_text_s=rec['terms']['first_text_s'],
                               after_speech_s=rec['terms']['after_speech_s'] + sec, llm_s=sec)
        if 'batch' in conds:
            text, sec = await batch(audio_dir / f'{m["id"]}.wav', args.lang)
            rec['batch'] = dict(text=text, first_text_s=None, after_speech_s=sec)
        print(f'  {m["id"]} 完成（{rec["seconds_of_audio"]} 秒录音）')
        return m['id'], rec

    print(f'{len(memos)} 条录音，语言 {args.lang}，实时模型 {args.model}，纠错模型 {args.llm}，'
          f'按真实语速推流，{datetime.now():%Y-%m-%d %H:%M}')
    results = dict(await asyncio.gather(*[one(m) for m in memos]))

    labels = {'plain': '识别时不给名单', 'terms': '识别时给 terms', 'correct': '识别后交给纠错模型', 'both': '识别时给 terms，再交给纠错模型',
              'batch': '（参考）整段上传识别，不给名单'}
    print(f'\n## 结果（{len(memos)} 条录音）\n')
    print('| 做法 | 人名写对 | 术语写对 | 开口后多久出第一段字（中位数） | 说完后多久拿到最终文字（中位数） |')
    print('|---|---|---|---|---|')
    summary = {}
    for c in conds:
        outs = {k: v[c]['text'] for k, v in results.items() if c in v}
        nh, nt, nmiss = score(outs, memos, names, args.lang)
        th, tt, _ = score(outs, memos, terms, args.lang)
        lat = [v[c]['after_speech_s'] for v in results.values() if c in v]
        firsts = [v[c]['first_text_s'] for v in results.values() if c in v and v[c]['first_text_s'] is not None]
        first_txt = '—（说完才开始）' if not firsts else f'{statistics.median(firsts):.1f} 秒'
        summary[c] = dict(names=f'{nh}/{nt}', terms=f'{th}/{tt}', name_misses=nmiss,
                          median_after_speech_s=statistics.median(lat), median_first_text_s=statistics.median(firsts) if firsts else None)
        print(f'| {labels[c]} | {nh}/{nt} | {th}/{tt} | {first_txt} | {statistics.median(lat):.2f} 秒 |')
    print('\n没写对的人名（做法: 名字×漏掉次数）：')
    for c in conds:
        print(f'- {labels[c]}：{summary[c]["name_misses"] or "无"}')

    Path(args.out).mkdir(exist_ok=True)
    path = Path(args.out) / f'{args.lang}_{args.model}_{datetime.now():%Y%m%d_%H%M}.json'
    path.write_text(json.dumps(dict(lang=args.lang, model=args.model, llm=args.llm, summary=summary,
                                    results=results), ensure_ascii=False, indent=1), encoding='utf-8')
    print(f'\n每条录音每种做法的文字和计时已存到 {path}')


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