"""5.8 Real-time speech recognition: should a proper-noun list go into the ASR model or a post-ASR correction LLM?

Streams the Module 2 sample voice memos (each with a verbatim ground-truth script) at real-time speaking speed to the AI Builder Space real-time speech endpoint, comparing:
  plain    No vocabulary list at the ASR stage
  terms    Pass the vocabulary list in ASR terms
  correct  No list during ASR; pass the raw transcript and vocabulary list to a post-ASR correction LLM (deepseek)
  both     Pass terms during ASR, then run post-ASR LLM correction
  batch    (Reference) Upload the full clip to the batch transcription endpoint after speech ends, with no list
Counts how many colleague names each configuration gets right and measures the post-stop wait for finalized text.

Usage:
    export AI_BUILDER_TOKEN=sk_...
    python run.py --lang zh            # Default: 23 name-bearing memos streamed at real-time pace (~10 min)
    python run.py --lang en
    python run.py --audio my_dir --memos my_memos.json   # Use your own recordings and ground truth; see README.md
"""
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          # Send 0.1 s of audio per frame: 24,000 samples/s * 2 bytes * 0.1 s

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('Please set the AI_BUILDER_TOKEN environment variable first (your AI Builder Space key)')
    return t


def ensure_audio(lang, audio_dir):
    """Download the 38 sample voice memos from Module 2 (~13 MB) on first run."""
    if audio_dir.exists() and any(audio_dir.glob('*.wav')):
        return
    print(f'Downloading sample audio to {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} is not a 24 kHz mono 16-bit WAV; convert it with ffmpeg first (see README.md)')
        return w.readframes(w.getnframes())


async def realtime(pcm, lang, model, terms):
    """Stream one recording. Returns (final_text, seconds_to_first_text, seconds_after_speech_to_final_text)."""
    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)              # Stream at real-time pace to simulate live speech
        t_end = time.time()                       # Speech finished
        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):
    """Batch transcription: upload the full file after speech ends. Returns (text, upload_to_result_seconds)."""
    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):
    """Count occurrences of a term in text. Chinese: case- and space-insensitive substring count; English: case-insensitive whole-word match."""
    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):
    """If a term appears n times in the ground-truth script and k times in the transcript, count min(n, k) hits."""
    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='Real-time speech model: gpt-live-transcribe or gpt-realtime')
    ap.add_argument('--llm', default='deepseek', help='LLM used for post-ASR correction')
    ap.add_argument('--memos', help='Ground-truth file (default: memos_<lang>.json)')
    ap.add_argument('--audio', help='Audio directory (default: audio_<lang>/; auto-downloads sample audio if missing)')
    ap.add_argument('--only', help='Run only these memo IDs, comma-separated, e.g. W01,W02')
    ap.add_argument('--conditions', default='plain,terms,correct,both,batch')
    ap.add_argument('--workers', type=int, default=2, help='Number of concurrent real-time sessions')
    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                                  # Use the exact same vocabulary list for ASR terms and LLM correction
    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)]   # Only benchmark memos that mention colleague 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"]} done ({rec["seconds_of_audio"]} s audio)')
        return m['id'], rec

    print(f'{len(memos)} memos, lang={args.lang}, realtime model={args.model}, correction LLM={args.llm}, '
          f'streaming at real-time pace, {datetime.now():%Y-%m-%d %H:%M}')
    results = dict(await asyncio.gather(*[one(m) for m in memos]))

    labels = {'plain': 'No vocabulary list', 'terms': 'ASR terms only', 'correct': 'Post-ASR LLM correction', 'both': 'ASR terms + post-ASR LLM correction',
              'batch': '(Reference) Batch upload, no list'}
    print(f'\n## Results ({len(memos)} memos)\n')
    print('| Configuration | Names right | Terms right | Time to first text (median) | Post-stop wait for final text (median) |')
    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 = '- (starts after speech)' if not firsts else f'{statistics.median(firsts):.1f} s'
        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} s |')
    print('\nMissed colleague names (configuration: {name: miss_count}):')
    for c in conds:
        print(f'- {labels[c]}: {summary[c]["name_misses"] or "none"}')

    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'\nPer-memo transcripts and timings for each configuration saved to {path}')


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