import socketio
import time
import base64
import os
from datetime import datetime
from typing import Dict, Callable, Any, Optional
from functools import partial
from threading import Event

class ASEngineClient:
    def __init__(self, server_url: str, namespace: str, id:str):
        """
        初始化客户端
        :param server_url: 服务器URL，例如 'https://asengine.net'
        :param namespace: 命名空间，例如 '/socket.io/7100eebf7a909e1ee3f4e06766f512f0358ecb6ac51d5434c17d31dd70890dc1'
        """
        self.server_url = server_url
        self.namespace = namespace
        self.id = id
        self.sid = None
        self.connected = False
        
        # 创建SocketIO客户端
        self.sio = socketio.Client()
        
        # 存储路由和对应的回调函数
        self.routes: Dict[str, Dict[str, Callable]] = {}
        
        self._on_connect_callback = None
        self._on_connect_callback_entry=None
        self._on_connect_callback_initData=None
        
        # ACK相关
        self.ack_events: Dict[str, Event] = {}  # 存储ack_id和对应的Event对象
        self._register_base_events()
        
    def _register_base_events(self):
        """注册基础事件处理函数"""
        # 连接事件
        @self.sio.on('connect', namespace=self.namespace)
        def on_connect():
            print("---------------ASEngineClient on_connect----------------")
            self.sid = self.sio.sid
            print(f"Connected to server with sid: {self.sid}")

            self.connected = True
            self.sio.emit('join', {'id': self.id}, namespace=self.namespace)
            if self._on_connect_callback is not None:
                self._on_connect_callback(
                    self,
                    self._on_connect_callback_entry,
                    init_data=self._on_connect_callback_initData
                )
        
        # 断开连接事件
        @self.sio.on('disconnect', namespace=self.namespace)
        def on_disconnect():
            self.connected = False
            self.sid = None
            print("Disconnected from server")
            
        @self.sio.on('error', namespace=self.namespace)
        def on_error(data):
            print(f"Received error: {data}")
            self.connected = False
            self.sid = None

        # 响应事件
        @self.sio.on('response', namespace=self.namespace)
        def on_response(data):
            print(f"Received response: {data}")
        
        # ACK事件处理
        @self.sio.on('ack', namespace=self.namespace)
        def on_ack(data):
            ack_id = data.get('ack_id')
            if ack_id and ack_id in self.ack_events:
                event = self.ack_events[ack_id]
                event.set()  # 触发事件，通知等待的线程
            
    def register_on_connect_entry(self, callback: Callable, entry: Any, init_data: Any)->bool:
        """注册连接事件回调函数
        
        参数:
        - callback: 连接事件回调函数，接受 **kwargs 参数
        - entry: 回调关联的条目对象
        - init_data: 初始化数据
        
        回调关键字参数:
        - entry: 回调关联的条目对象
        - init_data: 初始化数据
        """
        if not callable(callback):
            return False
        self._on_connect_callback = callback
        self._on_connect_callback_entry=entry
        self._on_connect_callback_initData=init_data
        return True


    def register_route_with_entry(self, route: str,
                                callback: Optional[Callable] = None,
                                entry: Any = None,
                                init_data: Optional[Any] = None) -> bool:
        """注册带条目的路由回调函数
        
        参数:
        - route: 路由名称
        - callback: 回调函数，接受 **kwargs 参数
        - entry: 回调关联的条目对象
        - init_data: 初始化数据
        
        回调关键字参数:
        - entry: 回调关联的条目对象
        - data: 接收到的消息数据
        - init_data: 初始化数据
        """
        if not route:
            print("Route cannot be empty")
            return False

        if callback is None:
            bound_cb = self._default_callback
        else:bound_cb = callback
        self.routes[route] = {
            'callback': bound_cb,
            'entry': entry,
            'init_data': init_data,
        }

        @self.sio.on(route, namespace=self.namespace)
        def on_route_message(data, _route=route):
            if not self.connected: return
            route_info = self.routes.get(_route)
            if not route_info:
                return
            cb = route_info['callback']          # 已经绑定过 entry
            init = route_info.get('init_data')
            entry = route_info.get('entry')
            try:
                cb(self, entry, data=data, init_data=init)    # 传递关键字参数
            except Exception as e:
                print(f"Callback error on route '{_route}': {e}")

        return True

        
    def register_route(self, route: str, callback: Optional[Callable] = None, init_data: Optional[Any] = None):
        """注册路由回调函数
        
        参数:
        - route: 路由名称
        - callback: 回调函数，接受 **kwargs 参数
        - init_data: 初始化数据
        
        回调关键字参数:
        - entry: 回调关联的条目对象（此处为 None）
        - data: 接收到的消息数据
        - init_data: 初始化数据
        """
        if not route:
            print("Route cannot be empty")
            return False
        
        # 存储路由和回调
        self.routes[route] = {
            'callback': callback if callback else self._default_callback,
            'entry': None,
            'init_data': init_data
        }
        
        # 注册事件处理器
        @self.sio.on(route, namespace=self.namespace)
        def on_route_message(data):
            print(f"Received message on route '{route}': {data}")
            if 'callback' in self.routes[route]:
                try:
                    self.routes[route]['callback'](
                        self,
                        entry=None,
                        data=data,
                        init_data=init_data
                    )
                except Exception as e:
                    print(f"Error in callback for route '{route}': {e}")
            
        print(f"Route '{route}' registered successfully")
        return True
        
    def _default_callback(self, **kwargs):
        """默认回调函数
        
        关键字参数:
        - entry: 回调关联的条目对象
        - data: 接收到的消息数据
        - init_data: 初始化数据
        """
        data = kwargs.get('data', {})
        print(f"Default callback received: {data}")
        
    def connect(self, timeout: int = 5) -> bool:
        """
        连接到服务器
        :param timeout: 超时时间(秒)
        :return: 是否连接成功
        """
        try:
            self.sio.connect(
                self.server_url+"?id="+self.id,
                namespaces=[self.namespace],
                wait_timeout=timeout
            )
            # 等待连接建立
            start_time = time.time()
            while not self.connected and (time.time() - start_time) < timeout:
                print(f"Connecting To {self.server_url}")
                time.sleep(0.2)
            return self.connected
        except Exception as e:
            print(f"Connection failed: {e}")
            return False
        
    def disconnect(self):
        """断开与服务器的连接"""
        if self.connected:
            self.sio.disconnect()
            self.connected = False
            self.sid = None
            print("Disconnected from server")
        
    def send_text(self, route: str, text: str, id:str="", wait_for_ack: bool = False, timeout: int = 5, return_event: bool = False) -> Any:
        """
        发送文本消息
        :param route: 目标路由
        :param text: 要发送的文本
        :param id: 发送者ID
        :param wait_for_ack: 是否等待服务器ACK确认
        :param timeout: 等待ACK的超时时间（秒）
        :param return_event: 是否返回事件对象和ack_id而不是等待ACK
        :return: 如果wait_for_ack为True且return_event为False，返回是否收到ACK；如果return_event为True，返回(event, ack_id)；否则返回是否发送成功
        """
        if not self.connected:
            print("Not connected to server")
            return False
        if id == "":
            id = self.id

        try:
            # 生成唯一ack_id
            ack_id = f"{id}_{datetime.now().strftime('%Y%m%d%H%M%S%f')}"
            data = {
                'type': 'text', 
                'data': text, 
                'id': id,
                'ack_id': ack_id
            }
            self.sio.emit(route, data, namespace=self.namespace)
            if not wait_for_ack:
                return True
            
            # 初始化ACK事件
            event = Event()
            self.ack_events[ack_id] = event
            
            # 如果需要返回事件对象，直接返回
            if return_event:
                return event, ack_id
            
            # 等待ACK或超时
            received_ack = event.wait(timeout=timeout)
            
            # 清理ACK事件
            del self.ack_events[ack_id]
            
            if received_ack:
                return True
            else:
                print(f"ACK timeout for message: {ack_id}")
                return False
        except Exception as e:
            print(f"Failed to send text to route '{route}': {e}")
            return False
        
    def send_voice(self, route: str, voice_data: bytes, id:str=None, wait_for_ack: bool = False, timeout: int = 5) -> bool:
        """
        发送语音数据
        :param route: 目标路由
        :param voice_data: 语音数据字节流
        :param id: 发送者ID
        :param wait_for_ack: 是否等待服务器ACK确认
        :param timeout: 等待ACK的超时时间（秒）
        :return: 是否发送成功（若wait_for_ack为True，则返回是否收到ACK）
        """
        if id is None:
            id = self.id
        if not self.connected:
            print("Not connected to server")
            return False
        
        try:
            # 生成唯一ack_id
            ack_id = f"{id}_{datetime.now().strftime('%Y%m%d%H%M%S%f')}"
            
            # 发送数据
            self.sio.emit(
                route,
                {
                    'type': 'audio',
                    'data': voice_data,
                    'timestamp': int(time.time() * 1000),
                    'sampleRate': 16000,
                    'id': id,
                    'ack_id': ack_id
                },
                namespace=self.namespace
            )
            # 如果不需要等待ACK，直接返回True
            if not wait_for_ack:
                return True
            
            # 初始化ACK事件
            event = Event()
            self.ack_events[ack_id] = event
            
            # 等待ACK或超时
            received_ack = event.wait(timeout=timeout)
            
            # 清理ACK事件
            del self.ack_events[ack_id]
            
            if received_ack:
                return True
            else:
                print(f"ACK timeout for message: {ack_id}")
                return False
        except Exception as e:
            print(f"Failed to send voice to route '{route}': {e}")
            return False
        
    def send_image(self, route: str, file_path: str, id:str, wait_for_ack: bool = False, timeout: int = 5) -> bool:
        """
        发送图片文件
        :param route: 目标路由
        :param file_path: 图片文件路径
        :param id: 发送者ID
        :param wait_for_ack: 是否等待服务器ACK确认
        :param timeout: 等待ACK的超时时间（秒）
        :return: 是否发送成功（若wait_for_ack为True，则返回是否收到ACK）
        """
        if not self.connected:
            print("Not connected to server")
            return False
        
        
        try:
            # 检查文件是否存在
            if not os.path.exists(file_path):
                print(f"File not found: {file_path}")
                return False
            
            # 读取文件并进行base64编码
            with open(file_path, 'rb') as f:
                image_data = 'data:image/jpeg;base64,' + base64.b64encode(f.read()).decode('utf-8')

            # 获取文件名
            filename = os.path.basename(file_path)
            
            # 生成唯一ack_id
            ack_id = f"{id}_{datetime.now().strftime('%Y%m%d%H%M%S%f')}"
            
            # 发送数据
            self.sio.emit(
                route,
                {
                    'type': 'image',
                    'data': image_data,
                    'id': id,
                    'ack_id': ack_id
                },
                namespace=self.namespace
            )
            
            # 如果不需要等待ACK，直接返回True
            if not wait_for_ack:
                return True
            
            # 初始化ACK事件
            event = Event()
            self.ack_events[ack_id] = event
            
            # 等待ACK或超时
            received_ack = event.wait(timeout=timeout)
            
            # 清理ACK事件
            del self.ack_events[ack_id]
            
            if received_ack:
                return True
            else:
                print(f"ACK timeout for message: {ack_id}")
                return False
        except Exception as e:
            print(f"Failed to send image to route '{route}': {e}")
            return False
        
    def run_forever(self):
        """
        保持客户端运行"""
        if self.connected:
            try:
                while True:
                    time.sleep(1)
            except KeyboardInterrupt:
                print("Client stopped by user")
                self.disconnect()
        else:
            print("Not connected to server")


class HSAEngineClient(ASEngineClient):
    def __init__(self, server_url: str, namespace: str, id:str, chatmanager):
        super().__init__(server_url, namespace, id)
        self.chatmanager = chatmanager
    def loadmarkdown(self, fmd, fmdp, fsp):
        event1, ack_id1 = self.send_text('markdown-in', fmd, self.id, wait_for_ack=True, return_event=True)
        event2, ack_id2 = self.send_text('markdown-prompt-in', fmdp, self.id, wait_for_ack=True, return_event=True)
        event3, ack_id3 = self.send_text('score-prompt-in', fsp, self.id, wait_for_ack=True, return_event=True)
        timeout = 5  # 超时时间
        success1 = event1.wait(timeout=timeout)
        success2 = event2.wait(timeout=timeout)
        success3 = event3.wait(timeout=timeout)
        if ack_id1 in self.ack_events:
            del self.ack_events[ack_id1]
        if ack_id2 in self.ack_events:
            del self.ack_events[ack_id2]
        if ack_id3 in self.ack_events:
            del self.ack_events[ack_id3]
        return success1 and success2 and success3