# funasr_client_api.py
import websocket
import threading
import time
import json
import ssl
from typing import Optional


class Funasr_websocket_recognizer:
    def __init__(self, host: str, port: str, is_ssl: bool = False, mode: str = "offline",
                 chunk_size: str = "0,10,5", ping_interval: int = 30):
        """
        FunASR WebSocket 识别器初始化
        :param host: FunASR 服务端IP
        :param port: FunASR 服务端端口
        :param is_ssl: 是否使用SSL（wss）
        :param mode: 识别模式（offline/online）
        :param chunk_size: 分片参数（仅online模式有效）
        :param ping_interval: 心跳间隔（秒）
        """
        self.host = host
        self.port = port
        self.is_ssl = is_ssl
        self.mode = mode
        self.chunk_size = chunk_size
        self.ping_interval = ping_interval

        # WebSocket 连接配置
        self.ws_url = f"{'wss' if is_ssl else 'ws'}://{host}:{port}/websocket/asr"
        self.ws = None
        self.ping_thread = None
        self.stop_ping = False
        self.result_buffer = []
        self.lock = threading.Lock()

    def _connect(self):
        """建立 WebSocket 连接"""
        ssl_opt = {"cert_reqs": ssl.CERT_NONE} if self.is_ssl else None
        self.ws = websocket.create_connection(self.ws_url, sslopt=ssl_opt)

        # 发送初始化消息
        init_msg = {
            "mode": self.mode,
            "chunk_size": self.chunk_size,
            "format": "wav",
            "rate": 16000,
            "need_pun": True,  # 是否需要标点
            "encoding": "raw"
        }
        self.ws.send(json.dumps(init_msg))

        # 启动心跳线程
        self.stop_ping = False
        self.ping_thread = threading.Thread(target=self._ping_worker)
        self.ping_thread.daemon = True
        self.ping_thread.start()

    def _ping_worker(self):
        """心跳保活"""
        while not self.stop_ping:
            try:
                if self.ws:
                    self.ws.ping()
                time.sleep(self.ping_interval)
            except Exception as e:
                print(f"Ping 失败: {e}")
                break

    def feed_chunk(self, audio_data: bytes, wait_time: float = 0.02) -> Optional[str]:
        """
        发送音频分片并获取识别结果
        :param audio_data: 音频字节数据（16k 16bit 单声道）
        :param wait_time: 等待结果的时间（秒）
        :return: 实时识别文本（None 表示无结果）
        """
        # 首次调用时建立连接
        if not self.ws:
            self._connect()

        try:
            # 发送音频数据（Base64 编码）
            import base64
            audio_b64 = base64.b64encode(audio_data).decode("utf-8")
            data_msg = {
                "data": audio_b64,
                "is_last": False  # 标记为非最后一包
            }
            self.ws.send(json.dumps(data_msg))

            # 等待并读取结果
            time.sleep(wait_time)
            result = None
            with self.lock:
                while self.ws and self.ws.connected:
                    if self.ws.readyState == websocket.ABNF.READYSTATE_CLOSING:
                        break
                    try:
                        resp = self.ws.recv(timeout=0.01)
                        if resp:
                            resp_json = json.loads(resp)
                            if "text" in resp_json:
                                result = resp_json["text"]
                                self.result_buffer.append(result)
                            break
                    except websocket.WebSocketTimeoutException:
                        break
            return result
        except Exception as e:
            print(f"发送分片失败: {e}")
            return None

    def close(self, timeout: float = 3) -> str:
        """
        关闭连接并获取最终识别结果
        :param timeout: 等待最终结果的超时时间（秒）
        :return: 最终识别文本
        """
        final_text = ""
        try:
            # 发送最后一包数据标记
            if self.ws and self.ws.connected:
                last_msg = {
                    "data": "",
                    "is_last": True
                }
                self.ws.send(json.dumps(last_msg))

                # 等待最终结果
                end_time = time.time() + timeout
                while time.time() < end_time:
                    try:
                        resp = self.ws.recv(timeout=0.1)
                        if resp:
                            resp_json = json.loads(resp)
                            if "text" in resp_json:
                                final_text = resp_json["text"]
                            if resp_json.get("is_final", False):
                                break
                    except websocket.WebSocketTimeoutException:
                        continue

            # 停止心跳并关闭连接
            self.stop_ping = True
            if self.ping_thread:
                self.ping_thread.join(timeout=1)
            if self.ws:
                self.ws.close()
        except Exception as e:
            print(f"关闭连接失败: {e}")

        # 合并所有结果
        with self.lock:
            all_text = "".join(self.result_buffer + [final_text])
        return all_text