# -*- coding: utf-8 -*-
"""
AIfeng/2025-01-27
FunASR语音识别模块
基于BaseASR的FunASR WebSocket客户端实现
"""

import json
import time
import asyncio
import websockets
import threading
import numpy as np
from threading import Thread, Event
from typing import Optional, Callable
import queue

from baseasr import BaseASR
import config_util as cfg
import util

class FunASRClient(BaseASR):
    """FunASR WebSocket客户端"""
    
    def __init__(self, opt, parent=None):
        # 确保opt有必要的属性
        if not hasattr(opt, 'fps'):
            opt.fps = 50  # 默认50fps
        if not hasattr(opt, 'batch_size'):
            opt.batch_size = 1
        if not hasattr(opt, 'l'):
            opt.l = 10
        if not hasattr(opt, 'r'):
            opt.r = 10
        
        super().__init__(opt, parent)
        
        # FunASR配置
        self.server_url = f"ws://{cfg.local_asr_ip}:{cfg.local_asr_port}"
        self.username = getattr(opt, 'username', 'default_user')
        
        # 连接状态
        self.websocket = None
        self.connected = False
        self.running = False
        self.reconnect_delay = getattr(cfg, 'asr_reconnect_delay', 1)
        self.max_reconnect_attempts = getattr(cfg, 'asr_max_reconnect_attempts', 5)
        
        # 消息队列
        self.message_queue = queue.Queue()
        self.result_queue = queue.Queue()
        
        # 线程控制
        self.connection_thread = None
        self.message_thread = None
        self.stop_event = Event()
        
        # 回调函数
        self.on_result_callback = None
        
        util.log(1, f"FunASR客户端初始化完成，服务器: {self.server_url}")
    
    def set_result_callback(self, callback: Callable[[str], None]):
        """设置识别结果回调函数
        
        Args:
            callback: 回调函数，接收识别结果字符串
        """
        self.on_result_callback = callback
    
    async def _connect_websocket(self):
        """连接WebSocket服务器"""
        try:
            # 修复: websockets新版本不支持timeout参数，使用asyncio.wait_for包装
            timeout_seconds = getattr(cfg, 'asr_timeout', 30)
            self.websocket = await asyncio.wait_for(
                websockets.connect(self.server_url),
                timeout=timeout_seconds
            )
            self.connected = True
            util.log(1, f"FunASR WebSocket连接成功: {self.server_url}")
            
            # 发送初始化配置消息（参考funasr_client_api.py）
            await self._send_init_message()
            
            return True
        except Exception as e:
            util.log(3, f"FunASR WebSocket连接失败: {e}")
            self.connected = False
            return False
    
    async def _send_init_message(self):
        """发送FunASR初始化配置消息"""
        try:
            # 根据参考项目funasr_client_api.py的格式
            init_message = {
                "mode": "2pass",
                "chunk_size": [0, 10, 5],  # [vad_need, chunk_size, chunk_interval]
                "encoder_chunk_look_back": 4,
                "decoder_chunk_look_back": 1, 
                "chunk_interval": 10,
                "wav_name": self.username,
                "is_speaking": True
            }
            
            await self.websocket.send(json.dumps(init_message))
            util.log(1, f"发送FunASR初始化消息: {init_message}")
            
        except Exception as e:
            util.log(3, f"发送初始化消息失败: {e}")
            raise e
    
    async def _disconnect_websocket(self):
        """断开WebSocket连接"""
        if self.websocket:
            try:
                await self.websocket.close()
            except Exception as e:
                util.log(2, f"关闭WebSocket连接时出错: {e}")
            finally:
                self.websocket = None
                self.connected = False
    
    async def _send_message(self, message: dict):
        """发送消息到FunASR服务器
        
        Args:
            message: 要发送的消息字典
        """
        if not self.connected or not self.websocket:
            util.log(2, "WebSocket未连接，无法发送消息")
            return False
        
        try:
            await self.websocket.send(json.dumps(message))
            return True
        except Exception as e:
            util.log(3, f"发送消息失败: {e}")
            self.connected = False
            return False
    
    async def _receive_messages(self):
        """接收WebSocket消息"""
        while self.connected and self.websocket:
            try:
                message = await asyncio.wait_for(
                    self.websocket.recv(), 
                    timeout=1.0
                )
                self._handle_recognition_result(message)
            except asyncio.TimeoutError:
                continue
            except websockets.exceptions.ConnectionClosed:
                util.log(2, "WebSocket连接已关闭")
                self.connected = False
                break
            except Exception as e:
                util.log(3, f"接收消息时出错: {e}")
                self.connected = False
                break
    
    async def _send_message_loop(self):
        """发送消息循环"""
        while self.connected and self.websocket:
            try:
                # 检查消息队列
                try:
                    message = self.message_queue.get_nowait()
                    
                    if isinstance(message, dict):
                        # JSON消息（配置消息或结束信号）
                        await self.websocket.send(json.dumps(message))
                        util.log(1, f"发送JSON消息: {message}")
                    elif isinstance(message, bytes):
                        # 二进制音频数据（参考funasr_client_api.py的feed_chunk方法）
                        # 确保音频数据以二进制格式发送
                        await self.websocket.send(message)
                        util.log(1, f"发送音频数据: {len(message)} bytes")
                    else:
                        util.log(2, f"未知消息类型: {type(message)}")
                        
                except queue.Empty:
                    # 队列为空，短暂等待
                    await asyncio.sleep(0.01)
                    
            except websockets.exceptions.ConnectionClosed:
                util.log(2, "发送消息时连接已关闭")
                self.connected = False
                break
            except Exception as e:
                util.log(3, f"发送消息时出错: {e}")
                self.connected = False
                break
    
    def _handle_recognition_result(self, message: str):
        """处理识别结果
        
        Args:
            message: 识别结果消息
        """
        try:
            # 尝试解析JSON
            try:
                result_data = json.loads(message)
                if isinstance(result_data, dict) and 'text' in result_data:
                    recognized_text = result_data['text']
                else:
                    recognized_text = message
            except json.JSONDecodeError:
                recognized_text = message
            
            # 存储结果
            self.result_queue.put(recognized_text)
            
            # 调用回调函数
            if self.on_result_callback:
                self.on_result_callback(recognized_text)
            
            # 发送到WebSocket服务器（兼容原有逻辑）
            self._send_to_web_clients(recognized_text)
            
            util.log(1, f"识别结果: {recognized_text}")
            
        except Exception as e:
            util.log(3, f"处理识别结果时出错: {e}")
    
    def _send_to_web_clients(self, text: str):
        """发送识别结果到Web客户端
        
        Args:
            text: 识别文本
        """
        try:
            from core import get_web_instance, get_instance
            
            # 发送到Web客户端
            if get_web_instance().is_connected(self.username):
                get_web_instance().add_cmd({
                    "panelMsg": text, 
                    "Username": self.username
                })
            
            # Human客户端通知改为日志记录（避免重复通知当前服务）
            util.log(1, f"FunASR识别结果[{self.username}]: {text}")
                
        except Exception as e:
            util.log(2, f"发送到Web客户端失败: {e}")
    
    async def _connection_loop(self):
        """连接循环，处理重连逻辑"""
        reconnect_attempts = 0
        
        while self.running and not self.stop_event.is_set():
            if not self.connected:
                util.log(1, f"尝试连接FunASR服务器 (第{reconnect_attempts + 1}次)")
                
                if await self._connect_websocket():
                    reconnect_attempts = 0
                    # 启动消息处理任务
                    receive_task = asyncio.create_task(self._receive_messages())
                    send_task = asyncio.create_task(self._send_message_loop())
                    
                    # 等待任务完成或连接断开
                    try:
                        await asyncio.gather(receive_task, send_task)
                    except Exception as e:
                        util.log(3, f"连接任务异常: {e}")
                    finally:
                        receive_task.cancel()
                        send_task.cancel()
                else:
                    reconnect_attempts += 1
                    if reconnect_attempts >= self.max_reconnect_attempts:
                        util.log(3, f"达到最大重连次数({self.max_reconnect_attempts})，停止重连")
                        break
                    
                    # 等待后重连
                    await asyncio.sleep(self.reconnect_delay)
                    self.reconnect_delay = min(self.reconnect_delay * 2, 30)  # 指数退避
            else:
                await asyncio.sleep(0.1)
        
        await self._disconnect_websocket()
    
    def _run_async_loop(self):
        """在独立线程中运行异步事件循环"""
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        
        try:
            loop.run_until_complete(self._connection_loop())
        except Exception as e:
            util.log(3, f"异步循环出错: {e}")
        finally:
            loop.close()
    
    def start(self):
        """启动FunASR客户端"""
        if self.running:
            util.log(2, "FunASR客户端已在运行")
            return
        
        self.running = True
        self.stop_event.clear()
        
        # 启动连接线程
        self.connection_thread = Thread(target=self._run_async_loop, daemon=True)
        self.connection_thread.start()
        
        util.log(1, "FunASR客户端已启动")
    
    def stop(self):
        """停止FunASR客户端"""
        if not self.running:
            return
        
        util.log(1, "正在停止FunASR客户端...")
        
        self.running = False
        self.stop_event.set()
        
        # 等待线程结束
        if self.connection_thread and self.connection_thread.is_alive():
            self.connection_thread.join(timeout=5)
        
        util.log(1, "FunASR客户端已停止")
    
    def send_audio_file(self, file_path: str):
        """发送音频文件进行识别
        
        Args:
            file_path: 音频文件路径
        """
        if not self.connected:
            util.log(2, "WebSocket未连接，无法发送音频文件")
            return False
        
        message = {"url": file_path}
        
        # 将消息放入队列，由异步线程处理
        self.message_queue.put(message)
        return True
    
    def send_audio(self, audio_data: bytes):
        """发送音频数据进行识别
        
        Args:
            audio_data: 音频字节数据
        """
        if not self.connected:
            util.log(2, "WebSocket未连接，无法发送音频数据")
            return False
        
        # 将音频数据放入队列
        self.message_queue.put(audio_data)
        return True
    
    def send_end_signal(self):
        """发送结束信号"""
        if not self.connected:
            return
        
        try:
            # 发送结束消息（参考funasr_client_api.py的close方法）
            end_message = {"is_speaking": False}
            self.message_queue.put(end_message)
            util.log(1, "发送FunASR结束信号")
        except Exception as e:
            util.log(3, f"发送结束信号失败: {e}")
    
    def start_recognition(self):
        """开始语音识别"""
        if not self.connected:
            self.start()
        
        # 发送开始识别消息
        start_message = {
            'vad_need': False,
            'state': 'StartTranscription'
        }
        self.message_queue.put(start_message)
        util.log(1, "开始语音识别")
    
    def stop_recognition(self):
        """停止语音识别"""
        if not self.connected:
            return
        
        # 发送停止识别消息
        stop_message = {
            'vad_need': False,
            'state': 'StopTranscription'
        }
        self.message_queue.put(stop_message)
        util.log(1, "停止语音识别")
    
    def get_latest_result(self, timeout: float = 0.1) -> Optional[str]:
        """获取最新的识别结果
        
        Args:
            timeout: 超时时间
            
        Returns:
            识别结果字符串或None
        """
        try:
            return self.result_queue.get(timeout=timeout)
        except queue.Empty:
            return None
    
    def warm_up(self):
        """预热模型"""
        super().warm_up()
        self.start()
        
        # 等待连接建立
        max_wait = 10  # 最多等待10秒
        wait_time = 0
        while not self.connected and wait_time < max_wait:
            time.sleep(0.1)
            wait_time += 0.1
        
        if self.connected:
            util.log(1, "FunASR客户端预热完成")
        else:
            util.log(2, "FunASR客户端预热超时")
    
    def run_step(self):
        """运行一步处理"""
        # 处理待发送的消息
        try:
            while not self.message_queue.empty():
                message = self.message_queue.get_nowait()
                # 这里需要通过某种方式发送到异步线程
                # 简化实现：直接记录日志
                util.log(1, f"准备发送消息: {message}")
        except queue.Empty:
            pass
        
        # 调用父类方法
        super().run_step()
    
    def get_next_feat(self, block=True, timeout=None):
        """获取下一个特征
        
        Args:
            block: 是否阻塞
            timeout: 超时时间
            
        Returns:
            特征数据
        """
        # 简化实现，返回空特征
        return np.zeros((1, 50), dtype=np.float32)
    
    async def connect(self):
        """异步连接到FunASR服务器"""
        if self.connected:
            util.log(1, "FunASR客户端已连接")
            return True
        
        try:
            success = await self._connect_websocket()
            if success:
                # 启动消息处理任务
                self.receive_task = asyncio.create_task(self._receive_messages())
                self.send_task = asyncio.create_task(self._send_message_loop())
                util.log(1, "FunASR异步连接建立成功")
            return success
        except Exception as e:
            util.log(3, f"FunASR异步连接失败: {e}")
            return False
    
    async def disconnect(self):
        """异步断开连接"""
        try:
            # 取消任务
            if hasattr(self, 'receive_task'):
                self.receive_task.cancel()
            if hasattr(self, 'send_task'):
                self.send_task.cancel()
            
            # 断开WebSocket连接
            await self._disconnect_websocket()
            util.log(1, "FunASR异步连接已断开")
        except Exception as e:
            util.log(2, f"断开FunASR连接时出错: {e}")
    
    async def send_audio_data(self, audio_data):
        """异步发送音频数据"""
        try:
            if isinstance(audio_data, str):
                # Base64编码的音频数据，需要解码
                import base64
                audio_bytes = base64.b64decode(audio_data)
            elif isinstance(audio_data, bytes):
                audio_bytes = audio_data
            else:
                util.log(2, f"不支持的音频数据类型: {type(audio_data)}")
                return False
            
            # 参考funasr_client_api.py，音频数据需要按chunk发送
            # 计算stride（参考项目中的计算方式）
            chunk_interval = 10  # ms
            chunk_size = 10      # ms 
            stride = int(60 * chunk_size / chunk_interval / 1000 * 16000 * 2)
            
            # 如果音频数据较大，分块发送
            if len(audio_bytes) > stride:
                chunk_num = (len(audio_bytes) - 1) // stride + 1
                for i in range(chunk_num):
                    beg = i * stride
                    chunk_data = audio_bytes[beg:beg + stride]
                    self.message_queue.put(chunk_data)
                    util.log(1, f"发送音频块 {i+1}/{chunk_num}: {len(chunk_data)} bytes")
            else:
                # 小数据直接发送
                self.message_queue.put(audio_bytes)
                util.log(1, f"发送音频数据: {len(audio_bytes)} bytes")
            
            return True
        except Exception as e:
            util.log(3, f"发送音频数据失败: {e}")
            return False

    def __del__(self):
        """析构函数"""
        self.stop()

# 兼容性别名
FunASR = FunASRClient