import json
import base64
import os
import struct
import asyncio

import websockets
import time
from fastapi import FastAPI, WebSocket, Query
from starlette.websockets import WebSocketState

from app.config.env import env


def add_websocket_voice_recognise_route(app: FastAPI):
  @app.websocket("/ws_voice_recognise")
  async def ws_voice_recognise(receive_socket: WebSocket, user_id: str = Query(..., description="用户ID")):
    print("新用户连接", user_id)
    # 接收连接
    await receive_socket.accept()

    # 百炼WebSocket配置
    API_KEY = env.llm_key_bailian
    QWEN_MODEL = "qwen3-asr-flash-realtime"
    baseUrl = "wss://dashscope.aliyuncs.com/api-ws/v1/realtime"
    url = f"{baseUrl}?model={QWEN_MODEL}"

    headers = {
      "Authorization": f"Bearer {API_KEY}",
      "OpenAI-Beta": "realtime=v1"
    }

    bailian_socket = None

    try:
      # 连接百炼WebSocket
      bailian_socket = await websockets.connect(url, additional_headers=headers)
      print("已连接到百炼WebSocket")

      # 发送会话更新事件
      session_event = {
        "event_id": "event_123",
        "type": "session.update",
        "session": {
          "modalities": ["text"],
          "input_audio_format": "pcm",
          "sample_rate": 16000,
          "input_audio_transcription": {
            "language": "zh"
          },
          "turn_detection": {
            "type": "server_vad",
            "threshold": 0.2,
            "silence_duration_ms": 800
          }
        }
      }
      await bailian_socket.send(json.dumps(session_event))

      # 创建双向转发任务
      async def waiting_msg_from_front():
        """将前端PCM数据转发到百炼"""
        try:
          while True:
            # 接收前端PCM数据（二进制格式）
            pcm_data_string = await receive_socket.receive_text()
            pcm_data_list = json.loads(pcm_data_string)
            # print("pcm_data", pcm_data_list)
            # print(f"接收到PCM数组数据，长度: {len(pcm_data_list)}")

            # 将 int16 列表打包成 bytes
            pcm_bytes = struct.pack('<' + 'h' * len(pcm_data_list), *pcm_data_list)
            # 将二进制PCM数据转换为base64编码
            encoded_data = base64.b64encode(pcm_bytes).decode('utf-8')

            # 构造发送到百炼的事件
            audio_event = {
              "event_id": f"event_{int(time.time() * 1000)}",
              "type": "input_audio_buffer.append",
              "audio": encoded_data
            }
            await bailian_socket.send(json.dumps(audio_event))
            # print("已转发音频数据到百炼", f"PCM二进制数据，长度: {len(pcm_bytes)} 字节")

        except Exception as e:
          print(f"等待前端的消息出错: {e}")

      async def waiting_msg_from_bailian():
        """将百炼识别结果转发到客户端"""
        try:
          async for message in bailian_socket:
            data = json.loads(message)
            print("收到百炼消息", data.get("type"))

            # 将识别结果发送给前端
            await receive_socket.send_text(json.dumps(data, ensure_ascii=False))

            # 如果收到最终转录结果，可以单独处理
            if data.get("type") == "conversation.item.input_audio_transcription.completed":
              transcript = data.get("transcript", "")
              print(f"最终转录结果: {transcript}")
              # 可以发送专门的转录消息给前端
              await receive_socket.send_text(json.dumps({
                "type": "final_transcript",
                "text": transcript
              }, ensure_ascii=False))

        except Exception as e:
          print(f"等待百炼的消息出错: {e}")

      # 并发运行两个任务
      await asyncio.gather(
        waiting_msg_from_front(),
        waiting_msg_from_bailian(),
        return_exceptions=True
      )

    except Exception as e:
      print(f"WebSocket连接错误: {e}")
    finally:
      print("finally ===>>>>>>>>>>>>>>>>>>>>>>>>>>>")
      # 清理连接
      if bailian_socket and bailian_socket.open:
        await bailian_socket.close()

      if receive_socket.application_state != WebSocketState.DISCONNECTED:
        print("关闭websocket")
        await receive_socket.close()
      print("websocket已经关闭", user_id)