#!/usr/bin/env python3
"""
Qwen3-Omni 实时聊天客户端 (基于 vLLM-Omni /v1/realtime WebSocket)

功能:
  1. 麦克风录音 → 实时流式发送音频 chunk 到服务端
  2. 服务端流式返回音频 delta → 实时播放
  3. 实时显示转录文本 (transcription.delta)
  4. 支持多轮对话

依赖:
  pip install websockets pyaudio numpy

使用方法:
  1. 先在远程服务器启动 vLLM-Omni 服务:
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8904

  2. 本地端口映射 (在 WSL/Linux 终端):
     ssh -fN -o ServerAliveInterval=60 -o ServerAliveCountMax=3 \
         -L 8904:127.0.0.1:8904 minor123@semi-ai.tech -p 60007

  3. 运行客户端:
     python test_qwen3_omni_realtime.py

交互:
  - 按 Enter 开始录音，再按 Enter 停止并发送
  - 输入 'q' 退出
  - 输入 't' 切换到纯文本模式 (不录音，直接打字)
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import json
import os
import queue
import struct
import sys
import threading
import time
import wave
from pathlib import Path

# ---------------------------------------------------------------------------
# 可选依赖检查
# ---------------------------------------------------------------------------
try:
    import websockets
except ImportError:
    print("请安装 websockets: pip install websockets")
    raise SystemExit(1)

try:
    import pyaudio
except ImportError:
    print("请安装 pyaudio: pip install pyaudio")
    print("  Windows: pip install pipwin && pipwin install pyaudio")
    print("  或者从 https://www.lfd.uci.edu/~gohlke/pythonlibs/#pyaudio 下载whl安装")
    raise SystemExit(1)

try:
    import numpy as np
except ImportError:
    print("请安装 numpy: pip install numpy")
    raise SystemExit(1)


# ======================== 配置 ========================
SAMPLE_RATE_INPUT = 16000    # 输入采样率 (Qwen3-Omni 要求 16kHz)
SAMPLE_RATE_OUTPUT = 24000   # 输出采样率 (模型输出 24kHz)
CHANNELS = 1                 # 单声道
SAMPLE_WIDTH = 2             # 16-bit PCM
CHUNK_MS = 120               # 每个音频 chunk 的毫秒数
CHUNK_SAMPLES = int(SAMPLE_RATE_INPUT * CHUNK_MS / 1000)
FORMAT = pyaudio.paInt16

# 录音配置
SILENCE_THRESHOLD = 500      # VAD 静音阈值 (RMS)
SILENCE_DURATION_SEC = 1.5   # 连续静音多少秒后自动停止录音


# ======================== 音频工具 ========================

class AudioRecorder:
    """麦克风录音器，支持实时回调获取音频数据"""

    def __init__(self, sample_rate: int = SAMPLE_RATE_INPUT, channels: int = CHANNELS):
        self.sample_rate = sample_rate
        self.channels = channels
        self.audio = pyaudio.PyAudio()
        self._stream: pyaudio.Stream | None = None
        self._frames: list[bytes] = []
        self._recording = False
        self._lock = threading.Lock()

    def list_devices(self):
        """列出所有音频输入设备"""
        print("\n可用的音频输入设备:")
        for i in range(self.audio.get_device_count()):
            info = self.audio.get_device_info_by_index(i)
            if info["maxInputChannels"] > 0:
                print(f"  [{i}] {info['name']} (采样率: {int(info['defaultSampleRate'])}Hz)")
        print()

    def _callback(self, in_data, frame_count, time_info, status):
        with self._lock:
            if self._recording:
                self._frames.append(in_data)
        return (None, pyaudio.paContinue)

    def start(self, device_index: int | None = None):
        """开始录音"""
        with self._lock:
            self._frames.clear()
            self._recording = True

        self._stream = self.audio.open(
            format=FORMAT,
            channels=self.channels,
            rate=self.sample_rate,
            input=True,
            input_device_index=device_index,
            frames_per_buffer=CHUNK_SAMPLES,
            stream_callback=self._callback,
        )
        self._stream.start_stream()

    def stop(self) -> bytes:
        """停止录音，返回完整 PCM16 数据"""
        with self._lock:
            self._recording = False

        if self._stream:
            self._stream.stop_stream()
            self._stream.close()
            self._stream = None

        return b"".join(self._frames)

    def is_recording(self) -> bool:
        return self._recording

    def close(self):
        if self._stream:
            self._stream.stop_stream()
            self._stream.close()
        self.audio.terminate()


class AudioPlayer:
    """实时音频播放器，支持流式播放 PCM delta"""

    def __init__(self, sample_rate: int = SAMPLE_RATE_OUTPUT):
        self.sample_rate = sample_rate
        self.audio = pyaudio.PyAudio()
        self._queue: queue.Queue[bytes | None] = queue.Queue()
        self._stream: pyaudio.Stream | None = None
        self._playing = False
        self._thread: threading.Thread | None = None

    def start(self):
        """启动播放线程"""
        self._playing = True
        self._stream = self.audio.open(
            format=FORMAT,
            channels=CHANNELS,
            rate=self.sample_rate,
            output=True,
            frames_per_buffer=CHUNK_SAMPLES,
        )
        self._thread = threading.Thread(target=self._play_loop, daemon=True)
        self._thread.start()

    def _play_loop(self):
        """播放循环，从队列取数据播放"""
        while self._playing:
            try:
                data = self._queue.get(timeout=0.1)
                if data is None:  # 停止信号
                    break
                if self._stream and data:
                    self._stream.write(data)
            except queue.Empty:
                continue

    def feed(self, pcm_data: bytes):
        """喂入一段 PCM 数据"""
        self._queue.put(pcm_data)

    def stop(self, drain: bool = True):
        """停止播放"""
        if drain:
            # 等待队列清空
            while not self._queue.empty():
                time.sleep(0.05)
            time.sleep(0.3)  # 额外等待确保播放完毕
        self._playing = False
        self._queue.put(None)  # 发送停止信号
        if self._thread:
            self._thread.join(timeout=2.0)
        if self._stream:
            self._stream.stop_stream()
            self._stream.close()
            self._stream = None

    def close(self):
        self.stop(drain=False)
        self.audio.terminate()


def save_wav(path: Path, pcm_data: bytes, sample_rate: int):
    """保存 PCM16 数据为 WAV 文件"""
    path.parent.mkdir(parents=True, exist_ok=True)
    with wave.open(str(path), "wb") as wf:
        wf.setnchannels(CHANNELS)
        wf.setsampwidth(SAMPLE_WIDTH)
        wf.setframerate(sample_rate)
        wf.writeframes(pcm_data)


# ======================== WebSocket 实时聊天客户端 ========================

async def realtime_chat_session(
    ws_url: str,
    model: str,
    recorder: AudioRecorder,
    player: AudioPlayer,
    save_audio_dir: Path | None = None,
    input_device_index: int | None = None,
):
    """
    单次实时对话 session:
      录音 → 流式发送 → 流式接收音频 + 文本 → 播放
    """
    text_chunks: list[str] = []
    audio_parts: list[bytes] = []
    output_sr = SAMPLE_RATE_OUTPUT

    print("\n🎤 开始录音... (按 Enter 停止)", flush=True)

    # 启动录音
    recorder.start(device_index=input_device_index)

    # 在另一个线程等待用户按 Enter 停止
    loop = asyncio.get_running_loop()

    def wait_for_stop():
        input()
        return True

    # 并行：录音 + 等待停止信号
    stop_future = loop.run_in_executor(None, wait_for_stop)

    async with websockets.connect(ws_url, max_size=64 * 1024 * 1024) as ws:
        # Step 1: 发送 session.update 声明模型
        await ws.send(json.dumps({"type": "session.update", "model": model}))

        # Step 2: 发送 input_audio_buffer.commit (非最终)
        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))

        # Step 3: 流式发送音频 chunk (边录边发)
        chunk_bytes = CHUNK_SAMPLES * SAMPLE_WIDTH
        buf = b""
        bytes_per_ms = SAMPLE_RATE_INPUT * SAMPLE_WIDTH // 1000
        sent_bytes = 0

        print("   (流式发送中...)", end="", flush=True)

        async def stream_audio():
            nonlocal sent_bytes
            # 从 recorder 读取已录制的帧
            last_read = 0
            while recorder.is_recording():
                await asyncio.sleep(0.01)
                current_frames = recorder._frames
                while last_read < len(current_frames):
                    frame = current_frames[last_read]
                    last_read += 1
                    # 按 chunk_bytes 切分发送
                    nonlocal buf
                    buf += frame
                    while len(buf) >= chunk_bytes:
                        chunk = buf[:chunk_bytes]
                        buf = buf[chunk_bytes:]
                        await ws.send(json.dumps({
                            "type": "input_audio_buffer.append",
                            "audio": base64.b64encode(chunk).decode("utf-8"),
                        }))
                        sent_bytes += len(chunk)
                        await asyncio.sleep(0)  # yield to event loop
            # 发送剩余数据
            if buf:
                await ws.send(json.dumps({
                    "type": "input_audio_buffer.append",
                    "audio": base64.b64encode(buf).decode("utf-8"),
                }))
                sent_bytes += len(buf)

        # 启动流式发送任务
        stream_task = asyncio.create_task(stream_audio())

        # 等待用户按 Enter
        await stop_future

        # 停止录音
        all_pcm = recorder.stop()
        print(f"\r📝 录音完成，共 {len(all_pcm) / SAMPLE_RATE_INPUT / SAMPLE_WIDTH:.1f}s，发送 {sent_bytes / SAMPLE_WIDTH / SAMPLE_RATE_INPUT:.1f}s", flush=True)

        # 等待流式发送完成
        stream_task.cancel()
        try:
            await stream_task
        except asyncio.CancelledError:
            pass

        # Step 4: 最终 commit，关闭输入流
        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))

        # Step 5: 接收服务端事件
        print("\n🤖 AI 回复:\n", flush=True)

        player.start()

        async def handle_events():
            nonlocal output_sr
            while True:
                message = await ws.recv()
                if isinstance(message, bytes):
                    continue
                event = json.loads(message)
                etype = event.get("type")

                if etype == "session.created":
                    continue

                if etype == "response.audio.delta":
                    sr = event.get("sample_rate_hz")
                    if isinstance(sr, int) and sr > 0:
                        output_sr = sr
                    audio_b64 = event.get("audio", "")
                    if audio_b64:
                        pcm_delta = base64.b64decode(audio_b64)
                        audio_parts.append(pcm_delta)
                        player.feed(pcm_delta)  # 实时播放
                    continue

                if etype == "transcription.delta":
                    delta = event.get("delta", "")
                    if delta:
                        text_chunks.append(delta)
                        print(delta, end="", flush=True)
                    continue

                if etype == "transcription.done":
                    final_text = event.get("text", "") or "".join(text_chunks)
                    usage = event.get("usage")
                    if text_chunks:
                        print()
                    if usage:
                        print(f"[usage: {usage}]")
                    continue

                if etype == "response.audio.done":
                    break

                if etype == "error":
                    raise RuntimeError(f"服务端错误: {event}")

        await handle_events()
        player.stop(drain=True)

    # 保存音频
    if save_audio_dir and audio_parts:
        all_audio = b"".join(audio_parts)
        save_dir = Path(save_audio_dir)
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        save_wav(save_dir / f"qwen3_omni_output_{timestamp}.wav", all_audio, output_sr)
        save_wav(save_dir / f"qwen3_omni_input_{timestamp}.wav", all_pcm, SAMPLE_RATE_INPUT)
        print(f"💾 音频已保存到: {save_dir}")

    return "".join(text_chunks), all_pcm, b"".join(audio_parts) if audio_parts else b""


# ======================== 纯文本流式模式 ========================
async def text_chat_session(ws_url: str, model: str):
    """
    纯文本聊天：使用 OpenAI 兼容的 /v1/chat/completions 流式端点
    适用于不需要音频输入，只想测试文本+音频输出的场景
    """
    import aiohttp

    print("\n💬 文本模式 (输入消息，Enter 发送，输入 'q' 返回主菜单)\n", flush=True)

    text_input = input("你的消息: ").strip()
    if not text_input or text_input.lower() == 'q':
        return

    # 将 ws://.../v1/realtime 转换为 http://.../v1/chat/completions
    http_url = ws_url.replace("ws://", "http://").replace("wss://", "https://")
    http_url = http_url.rsplit("/v1/realtime", 1)[0] + "/v1/chat/completions"

    payload = {
        "model": model,
        "messages": [{"role": "user", "content": text_input}],
        "stream": True,
        "modalities": ["text"],  # 仅文本，不要音频
    }

    print("🤖 AI: ", end="", flush=True)

    async with aiohttp.ClientSession() as session:
        async with session.post(http_url, json=payload) as resp:
            if resp.status != 200:
                text = await resp.text()
                print(f"\n错误 ({resp.status}): {text}")
                return

            full_text = ""
            async for line in resp.content:
                line = line.decode("utf-8").strip()
                if line.startswith("data: "):
                    data = line[6:]
                    if data == "[DONE]":
                        break
                    try:
                        chunk = json.loads(data)
                        delta = chunk.get("choices", [{}])[0].get("delta", {})
                        content = delta.get("content", "")
                        if content:
                            print(content, end="", flush=True)
                            full_text += content
                    except json.JSONDecodeError:
                        pass
            print()
    return full_text


# ======================== 文件模式 (无需麦克风) ========================

async def file_chat_session(
    ws_url: str,
    model: str,
    input_wav: Path,
    output_wav: Path | None,
    output_text: Path | None,
):
    """
    从 WAV 文件读取音频，发送到服务端，接收流式响应并保存。
    不需要麦克风，适合快速测试连通性。

    WAV 格式要求: mono, 16-bit PCM, 16kHz
    """
    # 读取 WAV 文件
    print(f"\n📂 读取音频文件: {input_wav}")
    with wave.open(str(input_wav), "rb") as wf:
        nch = wf.getnchannels()
        sw = wf.getsampwidth()
        fr = wf.getframerate()
        if nch != 1:
            print(f"⚠️ 音频是 {nch} 声道，将只使用第一声道")
        if sw != 2:
            raise ValueError(f"需要 16-bit PCM，当前 sample width={sw}")
        if fr != SAMPLE_RATE_INPUT:
            print(f"⚠️ 采样率 {fr}Hz != {SAMPLE_RATE_INPUT}Hz，服务端可能拒绝")
        pcm16 = wf.readframes(wf.getnframes())

    duration_sec = len(pcm16) / (fr * sw)
    print(f"   时长: {duration_sec:.1f}s, 大小: {len(pcm16)} bytes")

    text_chunks: list[str] = []
    audio_parts: list[bytes] = []
    output_sr = SAMPLE_RATE_OUTPUT
    chunk_bytes = CHUNK_SAMPLES * SAMPLE_WIDTH

    player = AudioPlayer()
    player.start()

    async with websockets.connect(ws_url, max_size=64 * 1024 * 1024) as ws:
        # session.update
        await ws.send(json.dumps({"type": "session.update", "model": model}))

        # 非最终 commit
        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))

        # 流式分块发送
        print("   📤 发送音频 chunk...", end="", flush=True)
        for i in range(0, len(pcm16), chunk_bytes):
            chunk = pcm16[i : i + chunk_bytes]
            await ws.send(json.dumps({
                "type": "input_audio_buffer.append",
                "audio": base64.b64encode(chunk).decode("utf-8"),
            }))
        print(" 完成")

        # 最终 commit
        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))

        # 接收事件
        print("\n🤖 AI 回复:\n", flush=True)
        while True:
            message = await ws.recv()
            if isinstance(message, bytes):
                continue
            event = json.loads(message)
            etype = event.get("type")

            if etype == "session.created":
                continue

            if etype == "response.audio.delta":
                sr = event.get("sample_rate_hz")
                if isinstance(sr, int) and sr > 0:
                    output_sr = sr
                audio_b64 = event.get("audio", "")
                if audio_b64:
                    pcm_delta = base64.b64decode(audio_b64)
                    audio_parts.append(pcm_delta)
                    player.feed(pcm_delta)
                continue

            if etype == "transcription.delta":
                delta = event.get("delta", "")
                if delta:
                    text_chunks.append(delta)
                    print(delta, end="", flush=True)
                continue

            if etype == "transcription.done":
                final_text = event.get("text", "") or "".join(text_chunks)
                usage = event.get("usage")
                if text_chunks:
                    print()
                if usage:
                    print(f"[usage: {usage}]")
                continue

            if etype == "response.audio.done":
                break

            if etype == "error":
                raise RuntimeError(f"服务端错误: {event}")

    player.stop(drain=True)
    player.close()

    # 保存输出
    all_audio = b"".join(audio_parts)
    if output_wav and all_audio:
        save_wav(output_wav, all_audio, output_sr)
        print(f"💾 音频已保存: {output_wav}")

    final_text = "".join(text_chunks)
    if output_text and final_text:
        output_text.parent.mkdir(parents=True, exist_ok=True)
        output_text.write_text(final_text, encoding="utf-8")
        print(f"💾 文本已保存: {output_text}")

    return final_text, all_audio


# ======================== 主程序 ========================

async def main():
    parser = argparse.ArgumentParser(
        description="Qwen3-Omni 实时聊天客户端 (vLLM-Omni WebSocket)"
    )
    parser.add_argument(
        "--url",
        default="ws://localhost:8904/v1/realtime",
        help="vLLM-Omni WebSocket 地址 (默认: ws://localhost:8904/v1/realtime)",
    )
    parser.add_argument(
        "--model",
        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
        help="模型名称 (默认: Qwen/Qwen3-Omni-30B-A3B-Instruct)",
    )
    parser.add_argument("--list-devices", action="store_true", help="列出音频设备后退出")
    parser.add_argument("--device", type=int, default=None, help="音频输入设备索引")
    parser.add_argument("--save-dir", default="./qwen3_omni_recordings", help="保存音频的目录")
    parser.add_argument("--text-only", action="store_true", help="纯文本模式 (不使用音频)")

    # 文件模式 (无需麦克风，快速测试)
    parser.add_argument("--input-wav", type=Path, default=None,
                        help="输入 WAV 文件路径 (mono, 16kHz, PCM16)。使用此参数进入文件测试模式，无需麦克风")
    parser.add_argument("--output-wav", type=Path, default=None,
                        help="输出 WAV 文件路径 (默认: <input>_output.wav)")
    parser.add_argument("--output-text", type=Path, default=None,
                        help="输出文本文件路径")

    args = parser.parse_args()

    # --- 文件模式 (快速测试，无需麦克风) ---
    if args.input_wav:
        input_path = args.input_wav
        if not input_path.exists():
            print(f"❌ 文件不存在: {input_path}")
            sys.exit(1)

        output_wav = args.output_wav or input_path.with_name(f"{input_path.stem}_output.wav")
        output_text = args.output_text or input_path.with_name(f"{input_path.stem}_output.txt")

        print("=" * 60)
        print("  Qwen3-Omni 文件模式测试")
        print("=" * 60)
        print(f"  服务地址: {args.url}")
        print(f"  输入文件: {input_path}")
        print(f"  输出音频: {output_wav}")
        print(f"  输出文本: {output_text}")
        print("=" * 60)

        try:
            text, audio = await file_chat_session(
                ws_url=args.url,
                model=args.model,
                input_wav=input_path,
                output_wav=output_wav,
                output_text=output_text,
            )
            if text:
                print(f"\n✅ 完成! 转录文本: {text}")
        except Exception as e:
            print(f"\n❌ 测试失败: {e}")
            import traceback
            traceback.print_exc()
            sys.exit(1)
        return

    # --- 交互模式 ---
    recorder = AudioRecorder()

    if args.list_devices:
        recorder.list_devices()
        recorder.close()
        return

    player = AudioPlayer()
    round_num = 0

    print("=" * 60)
    print("  Qwen3-Omni 实时聊天客户端")
    print("=" * 60)
    print(f"  服务地址: {args.url}")
    print(f"  模型: {args.model}")
    print()
    print("  交互方式:")
    print("    [Enter] - 开始录音 / 停止录音并获取回复")
    print("    [t]     - 切换到纯文本模式 (单轮)")
    print("    [q]     - 退出")
    print("=" * 60)

    try:
        while True:
            choice = input("\n>>> 按 Enter 开始语音对话 (t=文本, q=退出): ").strip().lower()

            if choice == 'q':
                print("👋 再见!")
                break

            if choice == 't':
                print("\n--- 纯文本模式 ---")
                await text_chat_session(args.url, args.model)
                continue

            # 语音对话
            round_num += 1
            print(f"\n--- 第 {round_num} 轮语音对话 ---")

            try:
                text, input_audio, output_audio = await realtime_chat_session(
                    ws_url=args.url,
                    model=args.model,
                    recorder=recorder,
                    player=player,
                    save_audio_dir=args.save_dir,
                    input_device_index=args.device,
                )
                if text:
                    print(f"\n📝 最终转录文本: {text}")
            except Exception as e:
                print(f"\n❌ 对话出错: {e}")
                import traceback
                traceback.print_exc()

    except KeyboardInterrupt:
        print("\n\n👋 用户中断, 再见!")
    finally:
        recorder.close()
        player.close()


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