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
from flask import current_app

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():
            self.connected = True
            self.sid = self.sio.get_sid(namespace=self.namespace)
            print(f"Connected to server. SID: {self.sid}")
            if self._on_connect_callback is not None:
                self._on_connect_callback(self._on_connect_callback_entry, 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:
                print(f"Received ACK for message: {ack_id}")
                event = self.ack_events[ack_id]
                event.set()  # 触发事件，通知等待的线程
            
    def register_on_connect_entry(self, callback: Callable, entry: Any, init_data: Any)->bool:
        if not callable(callback):
            return False
        self._on_connect_callback = callback
        self._on_connect_callback_entry=entry
        self._on_connect_callback_initData=init_data


    def register_route_with_entry(self, route: str,
                                callback: Optional[Callable] = None,
                                entry: Any = None,
                                init_data: Optional[Any] = None) -> bool:
        if not route:
            print("Route cannot be empty")
            return False

        if callback is None:
            bound_cb = self._default_callback
        else:
            bound_cb = partial(callback, entry)  # 等价于 lambda data, init: callback(entry, data, init)

        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')
            try:
                cb(data, init)                   # 现在只传 data / init_data
            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):
        """
        注册一个路由及其回调函数
        :param route: 路由名称
        :param callback: 回调函数，接收消息数据作为参数
        :return: 是否注册成功
        """
        if not route:
            print("Route cannot be empty")
            return False
        
        # 存储路由和回调
        self.routes[route] = {
            'callback': callback if callback else self._default_callback
        }
        
        # 注册事件处理器
        @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'](None, 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, entry: Any, data: Any, init_data: Optional[Any] = None):
        """默认回调函数"""
        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) -> bool:
        """
        发送文本消息
        :param route: 目标路由
        :param text: 要发送的文本
        :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
        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)
            print(f"Sent text to route '{route}': {text}")
            
            # 如果不需要等待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 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
            )
            
            print(f"Sent voice file to route '{route}'")
            
            # 如果不需要等待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
            )
            print(f"Sent image file to route '{route}': {filename}")
            
            # 如果不需要等待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):
        self.send_text('markdown-in', fmd, self.id, wait_for_ack=True)
        current_app.logger.debug(f"Sent markdown: {fmd}")
        self.send_text('markdown-prompt-in', fmdp, self.id, wait_for_ack=True)
        current_app.logger.debug(f"Sent markdown-prompt: {fmdp}")
        self.send_text('score-prompt-in', fsp, self.id, wait_for_ack=True)
        current_app.logger.debug(f"Sent score-prompt: {fsp}")