#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""语音转写（中泰英自动识别）

用法:
  python3 transcribe.py <音频文件>                 # 自动识别语言，small 模型
  python3 transcribe.py <音频文件> --lang th       # 指定泰语
  python3 transcribe.py <音频文件> --model base    # 快速模式（约4倍速，精度略低）
  python3 transcribe.py <音频文件> --segments      # 输出带时间戳的逐句结果

支持格式: mp3 / m4a / wav / ogg / flac / mp4 / webm 等（ffmpeg 能解的都能转）
"""
import os, sys, time, argparse
os.environ.setdefault('HF_ENDPOINT', 'https://hf-mirror.com')
os.environ.setdefault('HF_HUB_DISABLE_XET', '1')

from faster_whisper import WhisperModel

_MODELS = {}
def get_model(name='small'):
    if name not in _MODELS:
        _MODELS[name] = WhisperModel(name, device='cpu', compute_type='int8')
    return _MODELS[name]

def transcribe(path, language=None, model='small'):
    t0 = time.time()
    m = get_model(model)
    segments, info = m.transcribe(path, language=language, vad_filter=True, beam_size=5)
    segs = list(segments)
    text = ''.join(s.text for s in segs).strip()
    return {
        'language': info.language,
        'language_prob': round(info.language_probability, 3),
        'duration': round(info.duration, 1),
        'elapsed': round(time.time() - t0, 1),
        'text': text,
        'segments': [{'start': round(s.start, 2), 'end': round(s.end, 2), 'text': s.text.strip()} for s in segs],
    }

if __name__ == '__main__':
    ap = argparse.ArgumentParser()
    ap.add_argument('audio')
    ap.add_argument('--lang', default=None, help='语言提示: zh / th / en（默认自动识别）')
    ap.add_argument('--model', default='small', choices=['base', 'small', 'medium'], help='模型（默认 small）')
    ap.add_argument('--segments', action='store_true', help='输出逐句时间戳')
    args = ap.parse_args()

    r = transcribe(args.audio, language=args.lang, model=args.model)
    print(f"[语言={r['language']} 置信度={r['language_prob']} 音频{r['duration']}s 耗时{r['elapsed']}s]")
    print('=' * 40)
    if args.segments:
        for s in r['segments']:
            print(f"[{s['start']:>6.1f} - {s['end']:>6.1f}] {s['text']}")
    else:
        print(r['text'])