# -*- 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:
            self.websocket = await websockets.connect(
                self.server_url,
                timeout=getattr(cfg, 'asr_timeout', 30)
            )
            self.connected = True
            util.log(1, f"FunASR WebSocket连接成功: {self.server_url}")
            return True
        except Exception as e:
            util.log(3, f"FunASR WebSocket连接失败: {e}")
            self.connected = False
            return False
    
    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):
                        # 二进制音频数据
                        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 wsa_server
            
            # 发送到Web客户端
            if wsa_server.get_web_instance().is_connected(self.username):
                wsa_server.get_web_instance().add_cmd({
                    "panelMsg": text, 
                    "Username": self.username
                })
            
            # 发送到Human客户端
            if wsa_server.get_instance().is_connected_human(self.username):
                content = {
                    'Topic': 'human', 
                    'Data': {'Key': 'log', 'Value': text}, 
                    'Username': self.username
                }
                wsa_server.get_instance().add_cmd(content)
                
        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 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)
    
    def __del__(self):
        """析构函数"""
        self.stop()

# 兼容性别名
FunASR = FunASRClient