#!/usr/bin/env python3
"""
SenseVoice int8 语音识别 Demo
================================
功能：麦克风录音 → Silero VAD 语音分段 → SenseVoice 识别（带标点）→ 实时显示
同时监控并输出 CPU / 内存占用

用法：
    python3 demo/sense_voice_demo.py

首次运行会自动下载模型（约 228MB），请耐心等待。

按 Ctrl+C 退出时会输出性能统计。
"""

import sys
import os
import time
import argparse
import threading
import queue
import signal
from pathlib import Path

import numpy as np

# ── ANSI 颜色 ─────────────────────────────────────────────────
C = {
    "reset":  "\033[0m",
    "bold":   "\033[1m",
    "dim":    "\033[2m",
    "green":  "\033[32m",
    "cyan":   "\033[36m",
    "yellow": "\033[33m",
    "red":    "\033[31m",
    "magenta":"\033[35m",
    "blue":   "\033[94m",
}

def color(s, c):
    return f"{c}{s}{C['reset']}"

# ── 模型下载 ─────────────────────────────────────────────────
MODEL_DIR = Path(__file__).parent / "models"
SENSEVOICE_URL = (
    "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/"
    "sherpa-onnx-sense-voice-zh-en-ja-ko-yue-int8-2024-07-17.tar.bz2"
)
SILERO_VAD_URL = (
    "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/"
    "silero_vad.onnx"
)

SENSEVOICE_DIR = MODEL_DIR / "sherpa-onnx-sense-voice-zh-en-ja-ko-yue-int8-2024-07-17"
SENSEVOICE_MODEL = SENSEVOICE_DIR / "model.int8.onnx"
SENSEVOICE_TOKENS = SENSEVOICE_DIR / "tokens.txt"
SILERO_VAD_MODEL = MODEL_DIR / "silero_vad.onnx"


def download_with_progress(url, dest):
    """下载文件并显示进度条"""
    import urllib.request
    import shutil

    dest.parent.mkdir(parents=True, exist_ok=True)

    print(f"\n  📥 下载 {Path(dest).name} ...")
    print(f"     {url}")

    try:
        with urllib.request.urlopen(url) as resp:
            total = int(resp.headers.get("Content-Length", 0))
            downloaded = 0
            block_size = 8192
            with open(dest, "wb") as f:
                while True:
                    chunk = resp.read(block_size)
                    if not chunk:
                        break
                    f.write(chunk)
                    downloaded += len(chunk)
                    if total > 0:
                        pct = downloaded / total * 100
                        bar_len = 30
                        filled = int(bar_len * downloaded / total)
                        bar = "█" * filled + "░" * (bar_len - filled)
                        size_mb = downloaded / (1024 * 1024)
                        total_mb = total / (1024 * 1024)
                        print(
                            f"\r     [{bar}] {pct:5.1f}%  {size_mb:.0f}/{total_mb:.0f}MB",
                            end="", flush=True,
                        )
            print()
    except Exception as e:
        if dest.exists():
            dest.unlink()
        raise e


def ensure_models():
    """确保模型文件存在，不存在则下载"""
    import tarfile

    MODEL_DIR.mkdir(parents=True, exist_ok=True)

    # SenseVoice
    if not SENSEVOICE_MODEL.exists():
        archive = MODEL_DIR / "sense_voice_int8.tar.bz2"
        if not archive.exists():
            download_with_progress(SENSEVOICE_URL, archive)

        print(f"  📦 解压模型...")
        with tarfile.open(archive, "r:bz2") as tar:
            tar.extractall(path=MODEL_DIR)
        archive.unlink()
        print(f"  ✅ SenseVoice 模型就绪 ({SENSEVOICE_MODEL.stat().st_size / 1024 / 1024:.0f}MB)")

    # Silero VAD
    if not SILERO_VAD_MODEL.exists():
        download_with_progress(SILERO_VAD_URL, SILERO_VAD_MODEL)
        print(f"  ✅ Silero VAD 模型就绪")

    print()


# ── 性能监控 ─────────────────────────────────────────────────
class PerfMonitor:
    """后台线程采集 CPU / 内存"""

    def __init__(self, interval=1.0):
        self.interval = interval
        self._stop = threading.Event()
        self._thread = None
        self._lock = threading.Lock()

        # 累计统计
        self.cpu_percent_samples = []
        self.rss_mb_samples = []
        self.last_cpu = 0.0
        self.last_rss_mb = 0.0

        try:
            import psutil
            self._psutil = psutil
            self._proc = psutil.Process()
            self._available = True
        except ImportError:
            self._available = False

    @property
    def available(self):
        return self._available

    def start(self):
        if not self._available:
            return
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self):
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=2)

    def _run(self):
        while not self._stop.is_set():
            try:
                cpu = self._proc.cpu_percent(interval=None)
                rss_mb = self._proc.memory_info().rss / 1024 / 1024
            except Exception:
                cpu, rss_mb = 0, 0
            with self._lock:
                self.last_cpu = cpu
                self.last_rss_mb = rss_mb
                self.cpu_percent_samples.append(cpu)
                self.rss_mb_samples.append(rss_mb)
            self._stop.wait(self.interval)

    def current(self):
        with self._lock:
            return self.last_cpu, self.last_rss_mb

    def summary(self):
        with self._lock:
            if not self.cpu_percent_samples:
                return "无数据"
            cpu_avg = np.mean(self.cpu_percent_samples)
            cpu_max = np.max(self.cpu_percent_samples)
            rss_avg = np.mean(self.rss_mb_samples)
            rss_max = np.max(self.rss_mb_samples)
            return {
                "cpu_avg": cpu_avg,
                "cpu_max": cpu_max,
                "rss_avg_mb": rss_avg,
                "rss_max_mb": rss_max,
                "samples": len(self.cpu_percent_samples),
            }


# ── 识别统计 ─────────────────────────────────────────────────
class ASRStats:
    """记录每次识别的性能"""

    def __init__(self):
        self._lock = threading.Lock()
        self.segments = []  # list of (audio_duration, asr_time, text)

    def record(self, audio_duration, asr_time, text):
        with self._lock:
            self.segments.append((audio_duration, asr_time, text))

    def summary(self):
        with self._lock:
            if not self.segments:
                return None
            total_audio = sum(s[0] for s in self.segments)
            total_asr = sum(s[1] for s in self.segments)
            total_text = sum(len(s[2]) for s in self.segments)
            rtf = total_asr / total_audio if total_audio > 0 else 0
            return {
                "segments": len(self.segments),
                "total_audio_s": total_audio,
                "total_asr_s": total_asr,
                "total_chars": total_text,
                "rtf": rtf,
                "speedup": 1 / rtf if rtf > 0 else float("inf"),
            }


# ── 主程序 ────────────────────────────────────────────────────
def main():
    parser = argparse.ArgumentParser(description="SenseVoice int8 语音识别 Demo")
    parser.add_argument("--device", type=int, default=None,
                        help="音频输入设备编号（不指定则用默认）")
    parser.add_argument("--list-devices", action="store_true",
                        help="列出所有音频设备后退出")
    parser.add_argument("--save-audio", type=str, default=None,
                        help="保存最后一次识别的音频到指定文件（用于调试）")
    parser.add_argument("--num-threads", type=int, default=2,
                        help="ONNX 推理线程数（默认 2）")
    parser.add_argument("--vad-threshold", type=float, default=0.5,
                        help="VAD 灵敏度 0~1（默认 0.5，越小越敏感）")
    parser.add_argument("--silence-duration", type=float, default=0.8,
                        help="判定为句尾的静音时长秒数（默认 0.8）")
    parser.add_argument("--max-speech", type=float, default=30,
                        help="单段最长语音秒数（默认 30）")
    args = parser.parse_args()

    # ── 列出设备 ──
    import sounddevice as sd
    if args.list_devices:
        print("\n🎤 可用音频设备：\n")
        devices = sd.query_devices()
        for i, d in enumerate(devices):
            in_ch = d["max_input_channels"]
            if in_ch > 0:
                print(f"  [{i}] {d['name']} (输入通道: {in_ch}, 默认采样率: {d['default_samplerate']:.0f})")
        print()
        return

    # ── 下载模型 ──
    print(color("\n🔧 检查模型文件...", C["cyan"]))
    ensure_models()

    # ── 初始化 sherpa-onnx ──
    print(color("🚀 初始化识别引擎...", C["cyan"]))
    import sherpa_onnx

    # 通过工厂方法创建 SenseVoice 识别器
    recognizer = sherpa_onnx.OfflineRecognizer.from_sense_voice(
        model=str(SENSEVOICE_MODEL),
        tokens=str(SENSEVOICE_TOKENS),
        num_threads=args.num_threads,
        provider="cpu",
        language="auto",
        use_itn=True,  # 启用标点 + 逆文本正则化
    )

    # Silero VAD 配置
    silero_vad_config = sherpa_onnx.SileroVadModelConfig(
        model=str(SILERO_VAD_MODEL),
        threshold=args.vad_threshold,
        min_silence_duration=args.silence_duration,
        min_speech_duration=0.25,
        max_speech_duration=args.max_speech,
    )

    vad_config = sherpa_onnx.VadModelConfig(
        silero_vad=silero_vad_config,
        sample_rate=16000,
        num_threads=1,
    )

    vad = sherpa_onnx.VoiceActivityDetector(vad_config, buffer_size_in_seconds=60)

    print(color("✅ 引擎就绪！", C["green"]))

    # ── 启动性能监控 ──
    perf = PerfMonitor(interval=1.0)
    perf.start()

    asr_stats = ASRStats()

    # ── 音频参数 ──
    SAMPLE_RATE = 16000
    BLOCK_SIZE = 1024  # 每次读取的采样数

    # ── 优雅退出 ──
    running = True

    def on_sigint(sig, frame):
        nonlocal running
        running = False
        print(f"\n{color('⏹  正在退出...', C['yellow'])}")

    signal.signal(signal.SIGINT, on_sigint)

    # ── 打印标题 ──
    print()
    print(color("╔════════════════════════════════════════════════╗", C["bold"]))
    print(color("║     🎙️  SenseVoice 语音识别 Demo              ║", C["bold"]))
    print(color("║     中英混说 · 自带标点 · 低资源运行            ║", C["bold"]))
    print(color("╚════════════════════════════════════════════════╝", C["bold"]))
    print()
    print(f"  {color('🎤', C['cyan'])} 对着麦克风说话，说完停顿即可看到结果")
    print(f"  {color('⏸', C['cyan'])} 按 {color('Ctrl+C', C['yellow'])} 退出并查看统计")
    print(f"  {color('📊', C['cyan'])} VAD 阈值: {args.vad_threshold}  |  静音判定: {args.silence_duration}s")
    print()

    # 状态指示
    STATUS_IDLE = f" {color('⏳', C['dim'])} 等待语音..."
    STATUS_LISTENING = f" {color('🔴', C['red'])} 正在听..."
    STATUS_PROCESSING = f" {color('⚙️', C['yellow'])} 识别中..."

    sys.stdout.write(STATUS_IDLE)
    sys.stdout.flush()

    def audio_callback(indata, frames, time_info, status):
        """麦克风回调：将音频送入 VAD"""
        if status:
            print(f"\n  ⚠️ 音频状态: {status}")
        # 转为 float32 并归一化到 [-1, 1]
        if indata.dtype == np.int16:
            samples = indata[:, 0].astype(np.float32) / 32768.0
        else:
            samples = indata[:, 0].astype(np.float32)
        vad.accept_waveform(samples)

    try:
        with sd.InputStream(
            samplerate=SAMPLE_RATE,
            device=args.device,
            channels=1,
            dtype="float32",
            blocksize=BLOCK_SIZE,
            callback=audio_callback,
        ):
            last_status = STATUS_IDLE
            segment_count = 0

            while running:
                time.sleep(0.05)  # 50ms 轮询

                # 检测是否有语音段结束
                while not vad.empty():
                    speech_segment = vad.front  # sherpa_onnx.SpeechSegment
                    samples = np.array(speech_segment.samples, dtype=np.float32)
                    audio_duration = len(samples) / SAMPLE_RATE

                    vad.pop()

                    # ── 识别 ──
                    sys.stdout.write(f"\r{STATUS_PROCESSING} ({audio_duration:.1f}s 音频)")
                    sys.stdout.flush()

                    t0 = time.perf_counter()
                    stream = recognizer.create_stream()
                    stream.accept_waveform(SAMPLE_RATE, samples)
                    recognizer.decode_stream(stream)
                    text = stream.result.text.strip()
                    t1 = time.perf_counter()

                    asr_time = t1 - t0
                    asr_stats.record(audio_duration, asr_time, text)

                    # ── 显示结果 ──
                    segment_count += 1
                    rtf = asr_time / audio_duration if audio_duration > 0 else 0
                    speed = 1 / rtf if rtf > 0 else float("inf")

                    # 清除状态行
                    sys.stdout.write("\r" + " " * 60 + "\r")

                    # 输出识别文本
                    prefix = color(f"[{segment_count}]", C["dim"])
                    speed_info = color(
                        f"({asr_time:.2f}s / {audio_duration:.1f}s, RTF={rtf:.3f}, {speed:.0f}x)",
                        C["dim"],
                    )
                    print(f"  {prefix} {color(text, C['green'])}")
                    print(f"       {speed_info}")

                    # 保存音频（调试用）
                    if args.save_audio and segment_count == 1:
                        import sherpa_onnx
                        sherpa_onnx.write_wave(
                            str(args.save_audio),
                            samples.tolist(),
                            SAMPLE_RATE,
                        )
                        print(f"       💾 音频已保存到 {args.save_audio}")

                    # 恢复状态
                    sys.stdout.write(STATUS_IDLE)
                    sys.stdout.flush()

                # 更新状态行
                current_status = STATUS_LISTENING if vad.is_speech_detected() else STATUS_IDLE
                if current_status != last_status:
                    sys.stdout.write(f"\r{current_status}")
                    sys.stdout.flush()
                    last_status = current_status

    except KeyboardInterrupt:
        pass
    finally:
        running = False
        perf.stop()

    # ── 输出统计 ──
    print("\n")
    print(color("╔════════════════════════════════════════════════╗", C["bold"]))
    print(color("║               📊 性能统计                      ║", C["bold"]))
    print(color("╚════════════════════════════════════════════════╝", C["bold"]))
    print()

    # ASR 统计
    s = asr_stats.summary()
    if s:
        print(color("  🎙️  识别统计", C["cyan"]))
        print(f"     识别段数:     {s['segments']}")
        print(f"     总音频时长:   {s['total_audio_s']:.1f}s")
        print(f"     总识别耗时:   {s['total_asr_s']:.2f}s")
        print(f"     总字符数:     {s['total_chars']}")
        print(f"     RTF:         {s['rtf']:.4f}")
        print(f"     处理速度:     {s['speedup']:.1f}x 实时")
        print()

    # 系统资源统计
    ps = perf.summary()
    if isinstance(ps, dict):
        print(color("  💻 系统资源", C["cyan"]))
        print(f"     CPU 平均:    {ps['cpu_avg']:.1f}%")
        print(f"     CPU 峰值:    {ps['cpu_max']:.1f}%")
        print(f"     内存平均:    {ps['rss_avg_mb']:.0f}MB")
        print(f"     内存峰值:    {ps['rss_max_mb']:.0f}MB")
        print(f"     采样点数:    {ps['samples']}")
        print()
    elif not perf.available:
        print(color("  ⚠️  未安装 psutil，无法采集 CPU/内存数据", C["yellow"]))
        print(color("     安装: pip3 install psutil", C["dim"]))
        print()


if __name__ == "__main__":
    main()