#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
语音识别WebSocket服务器入口文件
提供基于WebSocket的实时语音识别服务，支持离线和在线两种模式
"""

import asyncio
import json
import websockets
import ssl
import argparse
from config import parse_args
from models import load_models
from service import ASRService

# 全局变量，存储当前连接的WebSocket客户端
websocket_users = set()


async def ws_reset(websocket):
    """重置WebSocket连接状态并关闭连接"""
    print(f"重置WebSocket连接，当前连接数: {len(websocket_users)}")

    # 重置状态字典
    websocket.status_dict_asr_online["cache"] = {}
    websocket.status_dict_asr_online["is_final"] = True
    websocket.status_dict_vad["cache"] = {}
    websocket.status_dict_vad["is_final"] = True
    websocket.status_dict_punc["cache"] = {}

    # 关闭连接
    await websocket.close()


async def clear_websocket():
    """清理所有WebSocket连接"""
    for websocket in websocket_users:
        await ws_reset(websocket)
    websocket_users.clear()


async def ws_serve(websocket, path):
    """
    WebSocket服务主函数，处理客户端连接和消息

    参数:
        websocket: WebSocket连接对象
        path: 连接路径
    """
    frames = []  # 存储所有音频帧
    frames_asr = []  # 存储用于离线ASR的音频帧
    frames_asr_online = []  # 存储用于在线ASR的音频帧

    global websocket_users
    # await clear_websocket()  # 清理现有连接(目前注释掉，允许多客户端)

    # 添加到用户集合
    websocket_users.add(websocket)

    # 初始化连接状态
    websocket.status_dict_asr = {}
    websocket.status_dict_asr_online = {"cache": {}, "is_final": False}
    websocket.status_dict_vad = {"cache": {}, "is_final": False}
    websocket.status_dict_punc = {"cache": {}}
    websocket.chunk_interval = 10
    websocket.vad_pre_idx = 0
    websocket.is_speaking = True  # 默认用户正在说话

    # 语音检测状态
    speech_start = False
    speech_end_i = -1

    # 初始化配置
    websocket.wav_name = "microphone"
    websocket.mode = "2pass"  # 默认使用两阶段识别模式

    print("新用户已连接", flush=True)

    try:
        # 持续接收客户端消息
        async for message in websocket:
            # 处理JSON配置消息
            if isinstance(message, str):
                try:
                    messagejson = json.loads(message)

                    # 更新各种配置参数
                    if "is_speaking" in messagejson:
                        websocket.is_speaking = messagejson["is_speaking"]
                        websocket.status_dict_asr_online["is_final"] = (
                            not websocket.is_speaking
                        )
                    if "chunk_interval" in messagejson:
                        websocket.chunk_interval = messagejson["chunk_interval"]
                    if "wav_name" in messagejson:
                        websocket.wav_name = messagejson.get("wav_name")
                    if "chunk_size" in messagejson:
                        chunk_size = messagejson["chunk_size"]
                        if isinstance(chunk_size, str):
                            chunk_size = chunk_size.split(",")
                        websocket.status_dict_asr_online["chunk_size"] = [
                            int(x) for x in chunk_size
                        ]
                    if "encoder_chunk_look_back" in messagejson:
                        websocket.status_dict_asr_online["encoder_chunk_look_back"] = (
                            messagejson["encoder_chunk_look_back"]
                        )
                    if "decoder_chunk_look_back" in messagejson:
                        websocket.status_dict_asr_online["decoder_chunk_look_back"] = (
                            messagejson["decoder_chunk_look_back"]
                        )
                    if "hotword" in messagejson:
                        websocket.status_dict_asr["hotword"] = messagejson["hotwords"]
                    if "mode" in messagejson:
                        websocket.mode = messagejson["mode"]
                except json.JSONDecodeError:
                    print(f"无效的JSON消息: {message}")

            # 根据chunk_interval更新VAD的chunk_size
            websocket.status_dict_vad["chunk_size"] = int(
                websocket.status_dict_asr_online.get("chunk_size", [0, 10])[1]
                * 60
                / websocket.chunk_interval
            )

            # 处理音频数据
            if (
                len(frames_asr_online) > 0
                or len(frames_asr) >= 0
                or not isinstance(message, str)
            ):
                if not isinstance(message, str):  # 二进制音频数据
                    # 添加到帧缓冲区
                    frames.append(message)
                    duration_ms = len(message) // 32  # 计算音频时长
                    websocket.vad_pre_idx += duration_ms

                    # 处理在线ASR
                    frames_asr_online.append(message)
                    websocket.status_dict_asr_online["is_final"] = speech_end_i != -1

                    # 达到chunk_interval或最终帧时处理在线ASR
                    if (
                        len(frames_asr_online) % websocket.chunk_interval == 0
                        or websocket.status_dict_asr_online["is_final"]
                    ):
                        if websocket.mode == "2pass" or websocket.mode == "online":
                            audio_in = b"".join(frames_asr_online)
                            try:
                                await asr_service.async_asr_online(websocket, audio_in)
                            except Exception as e:
                                print(f"在线ASR处理错误: {e}")
                        frames_asr_online = []

                    # 如果检测到语音开始，收集帧用于离线ASR
                    if speech_start:
                        frames_asr.append(message)

                    # VAD处理 - 语音活动检测
                    try:
                        speech_start_i, speech_end_i = await asr_service.async_vad(
                            websocket, message
                        )
                    except Exception as e:
                        print(f"VAD处理错误: {e}")

                    # 检测到语音开始
                    if speech_start_i != -1:
                        speech_start = True
                        # 计算开始偏移并收集前面的帧
                        beg_bias = (
                            websocket.vad_pre_idx - speech_start_i
                        ) // duration_ms
                        frames_pre = (
                            frames[-beg_bias:] if beg_bias < len(frames) else frames
                        )
                        frames_asr = []
                        frames_asr.extend(frames_pre)

                # 处理离线ASR (语音结束或用户停止说话)
                if speech_end_i != -1 or not websocket.is_speaking:
                    if websocket.mode == "2pass" or websocket.mode == "offline":
                        audio_in = b"".join(frames_asr)
                        try:
                            await asr_service.async_asr(websocket, audio_in)
                        except Exception as e:
                            print(f"离线ASR处理错误: {e}")

                    # 重置状态
                    frames_asr = []
                    speech_start = False
                    frames_asr_online = []
                    websocket.status_dict_asr_online["cache"] = {}

                    # 如果用户停止说话，完全重置
                    if not websocket.is_speaking:
                        websocket.vad_pre_idx = 0
                        frames = []
                        websocket.status_dict_vad["cache"] = {}
                    else:
                        # 保留最近的帧用于下一轮处理
                        frames = frames[-20:]

    except websockets.ConnectionClosed:
        print(f"连接已关闭...", flush=True)
        await ws_reset(websocket)
        websocket_users.remove(websocket)
    except websockets.InvalidState:
        print("无效的WebSocket状态...")
    except Exception as e:
        print(f"发生异常: {e}")


def start_server(args, asr_service_instance):
    """
    启动WebSocket服务器

    参数:
        args: 命令行参数
        asr_service_instance: ASR服务实例
    """
    global asr_service
    asr_service = asr_service_instance

    # 配置SSL (如果提供了证书)
    if args.certfile and len(args.certfile) > 0:
        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ssl_context.load_cert_chain(args.certfile, keyfile=args.keyfile)

        start_server = websockets.serve(
            ws_serve,
            args.host,
            args.port,
            subprotocols=["binary"],
            ping_interval=None,
            ssl=ssl_context,
        )
    else:
        start_server = websockets.serve(
            ws_serve, args.host, args.port, subprotocols=["binary"], ping_interval=None
        )

    print(f"WebSocket服务器已启动 - 监听 {args.host}:{args.port}")

    # 启动事件循环
    asyncio.get_event_loop().run_until_complete(start_server)
    asyncio.get_event_loop().run_forever()


if __name__ == "__main__":
    # 解析命令行参数
    args = parse_args()

    # 加载模型
    print("正在加载模型...")
    models = load_models(args)
    print("模型加载完成！当前仅支持单个客户端同时连接！")

    # 创建ASR服务
    asr_service = ASRService(models)

    # 启动服务器
    start_server(args, asr_service)
