# -*- coding: utf-8 -*-
"""
高性能BLE客户端 - 优化版本
修复了以下问题:
1. 内存泄漏 - 每次连接创建新实例导致资源无法释放
2. 日志阻塞 - 实现异步日志队列
3. 通知失效 - 完善连接状态检查和清理
4. WebSocket保活 - 添加ping/pong机制
5. 错误处理 - 全面增强异常捕获和处理
"""
import sys
import asyncio
import websockets
from bleak import BleakScanner, BleakClient
import json
import base64
import threading
from collections import defaultdict
import socket
import time
import traceback

import platform

# 方法1：通过sys模块快速判断（推荐）
if sys.platform.startswith('win32'):
    import winrt.windows.foundation.collections  # noqa
    import winrt.windows.devices.bluetooth  # noqa
    import winrt.windows.devices.bluetooth.advertisement  # noq
    print("当前运行在Windows系统")
elif sys.platform.startswith('linux'):
    print("当前运行在Linux系统")
elif sys.platform.startswith('darwin'):
    print("当前运行在macOS系统")
# 高性能日志系统
write_lock = threading.Lock()
_log_queue = None
_log_task = None
_log_dropped_count = 0

def log_message_sync(direction, message):
    """同步日志记录函数"""
    try:
        log_entry = f"{direction}: {message}\n"
        print(log_entry, end='')  # 控制台仍然输出
        
        # 截断过长消息，避免写入过大的日志
        if len(log_entry) > 4000:
            log_entry = log_entry[:4000] + "... [truncated]\n"
        
        with write_lock:
            with open('b.log', 'a', encoding='utf-8') as f:
                f.write(log_entry)
    except Exception:
        pass  # 日志失败不应影响主流程

async def _log_writer_loop():
    """后台日志写入循环"""
    global _log_queue
    while True:
        try:
            item = await _log_queue.get()
            if item is None:
                break
            direction, message = item
            log_message_sync(direction, message)
        except asyncio.CancelledError:
            break
        except Exception as e:
            await log_exception_async(e, "日志写入循环")
            await asyncio.sleep(0.05)

async def log_message(direction, message):
    """异步日志记录，非阻塞"""
    global _log_queue, _log_task, _log_dropped_count
    
    # 首次调用时初始化
    if _log_queue is None:
        _log_queue = asyncio.Queue(maxsize=1000)
        _log_task = asyncio.create_task(_log_writer_loop())
    
    try:
        _log_queue.put_nowait((direction, message))
    except asyncio.QueueFull:
        _log_dropped_count += 1
        if _log_dropped_count % 100 == 1:
            print(f"警告: 日志队列已满，已丢弃 {_log_dropped_count} 条日志")

async def log_exception_async(exc, context=""):
    """异步记录异常到日志"""
    try:
        exc_type = type(exc).__name__
        exc_msg = str(exc)
        exc_traceback = traceback.format_exc()
        error_msg = f"异常[{context}]: {exc_type}: {exc_msg}\n堆栈:\n{exc_traceback}"
        await log_message("异常", error_msg)
    except Exception:
        pass  # 记录异常失败时pass，避免循环异常

def log_exception_sync_to_async(exc, context=""):
    """同步函数中将异常放入异步日志队列"""
    try:
        exc_type = type(exc).__name__
        exc_msg = str(exc)
        exc_traceback = traceback.format_exc()
        error_msg = f"异常[{context}]: {exc_type}: {exc_msg}\n堆栈:\n{exc_traceback}"
        
        global _log_queue, _log_task
        if _log_queue is not None:
            try:
                _log_queue.put_nowait(("异常", error_msg))
            except asyncio.QueueFull:
                pass  # 队列满时pass
    except Exception:
        pass  # 记录异常失败时pass

def is_port_in_use(port, host='localhost'):
    """检查端口是否被占用"""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        try:
            s.bind((host, port))
            return False
        except socket.error:
            return True

async def self_test(websocket):
    """发送自检测试消息"""
    test_message = json.dumps({
        "jsonrpc": "2.0",
        "method": "ping",
        "params": {"timestamp": int(time.time())},
        "id": "test"
    })
    await log_message("自测", test_message)
    await websocket.send(test_message)
    
    try:
        # 等待响应，设置超时
        response = await asyncio.wait_for(websocket.recv(), timeout=5.0)
        await log_message("自测响应", response)
        return True
    except asyncio.TimeoutError as e:
        await log_exception_async(e, "自测超时")
        pass
        return False
    except Exception as e:
        await log_exception_async(e, "自测")
        pass
        return False

class BLEClient:
    def __init__(self):
        self.target_device = None
        self.client = None
        self.services = []
        self.optional_services = []
        self.websocket = None
        self.notification_records = defaultdict(lambda: (None, 0.0))  # 特征ID: (最后消息, 时间戳)
        self._notification_callbacks = {}  # 存储通知回调，用于清理
        self._shutdown = False  # 添加关闭标志
        self._scanner = None  # 当前扫描器实例

    def on_disconnect(self, client):
        print("BLE连接断开，关闭WebSocket")
        # 设置关闭标志
        self._shutdown = True
        # 在事件循环中关闭WebSocket
        if self.websocket and not self.websocket.closed:
            try:
                # 尝试获取当前事件循环
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    loop.create_task(self.close_websocket())
            except Exception as e:
                # 如果没有事件循环，直接设置为None（延迟清理）
                log_exception_sync_to_async(e, "on_disconnect获取事件循环")
                pass

    async def close_websocket(self):
        """主动关闭WebSocket连接"""
        if self.websocket:
            try:
                if not self.websocket.closed:
                    # 使用1000状态码正常关闭
                    await self.websocket.close(code=1000, reason="BLE connection closed")
                    print("WebSocket连接已关闭")
            except Exception as e:
                await log_exception_async(e, "关闭WebSocket")
                pass
            finally:
                self.websocket = None
                self._shutdown = True

    def detection_callback(self, device, advertisement_data):
        if any(service_uuid in advertisement_data.service_uuids for service_uuid in self.services):
            # 优先使用广播包中的local_name
            device_name = advertisement_data.local_name if advertisement_data.local_name else (device.name if device.name else "")
            
            # 丢弃名字为空的设备
            if not device_name or device_name.strip() == "":
                print(f"跳过名字为空的设备: {device.address}")
                return None
            
            self.target_device = (device, advertisement_data)
            if not self.target_device:
                print("未找到匹配设备")
                return
            else:
                device, adv_data = self.target_device
                print("\n找到目标设备:")
                print(f"设备名称(从广播包): {device_name}")
                print(f"设备地址: {device.address}")
                print(f"信号强度: {device.rssi} dBm")
                print("\n广播信息:")
                print(f"服务UUID列表: {adv_data.service_uuids}")
                print(f"制造商数据: {adv_data.manufacturer_data}")
                print(f"服务数据: {adv_data.service_data}")
                print(f"本地名称(local_name): {adv_data.local_name}")
                print(f"设备名称(device.name): {device.name if device.name else 'N/A'}")
                # 发现设备后，尝试立即停止扫描
                try:
                    if self._scanner is not None:
                        # 在事件循环中停止扫描，避免阻塞回调
                        asyncio.create_task(self._stop_scan_early())
                except Exception as e:
                    log_exception_sync_to_async(e, "detection_callback停止扫描")
                    pass
            return self.target_device

    async def _stop_scan_early(self):
        try:
            if self._scanner is not None:
                await self._scanner.stop()
        except Exception as e:
            await log_exception_async(e, "_stop_scan_early")
            pass

    async def handle_client(self, websocket, path):
        self.websocket = websocket
        self._shutdown = False  # 重置关闭标志
        if path != "/scratch/ble":
            await websocket.close(code=1003, reason="Path not allowed")
            return

        try:
            async for message in websocket:
                # 检查是否应该关闭连接
                if self._shutdown:
                    print("检测到关闭标志，终止消息接收")
                    break
                    
                try:
                    await log_message("接收", message)
                    request = json.loads(message)
                    
                    if request["jsonrpc"] != "2.0":
                        continue

                    method = request.get("method")
                    params = request.get("params", {})
                    request_id = request.get("id")

                    if method == "discover":
                        self.services = []
                        for filt in params.get("filters", [{}]):
                            self.services.extend(filt.get("services", []))
                        self.optional_services = params.get("optionalServices", [])

                        # 双重扫描：快速扫描后若未发现，再进行扩展扫描；发现即停
                        phases = [("active", 3.0), ("passive", 6.0)]
                        found = False
                        for phase_index, (scan_mode, duration) in enumerate(phases, start=1):
                            self.target_device = None
                            self._scanner = BleakScanner(scanning_mode=scan_mode)
                            self._scanner.register_detection_callback(self.detection_callback)
                            print(f"开始第{phase_index}阶段扫描（模式: {scan_mode}, 时长: {duration}s）...")
                            try:
                                await self._scanner.start()
                                # 轮询检查是否已找到，找到则提前停止
                                start_ts = time.time()
                                while time.time() - start_ts < duration and not self.target_device:
                                    await asyncio.sleep(0.1)
                            finally:
                                try:
                                    await self._scanner.stop()
                                except Exception as e:
                                    await log_exception_async(e, "扫描停止")
                                    pass
                                self._scanner = None
                            
                            if self.target_device:
                                found = True
                                break
                            else:
                                print(f"第{phase_index}阶段未找到设备")

                        if found:
                            device, adv_data = self.target_device
                            # 优先使用广播包中的local_name，如果没有则使用device.name
                            device_name = adv_data.local_name if adv_data.local_name else (device.name if device.name else "")
                            discover_response = json.dumps({
                                "jsonrpc": "2.0",
                                "method": "didDiscoverPeripheral",
                                "params": {
                                    "name": device_name,
                                    "peripheralId": device.address,
                                    "rssi": device.rssi
                                }
                            })
                            await log_message("下发", discover_response)
                            if not websocket.closed:
                                await websocket.send(discover_response)

                            result_response = json.dumps({
                                "jsonrpc": "2.0",
                                "result": None,
                                "id": request_id if request_id else 0
                            })
                            await log_message("下发", result_response)
                            if not websocket.closed:
                                await websocket.send(result_response)

                    elif method == "connect":
                        peripheral_id = params.get("peripheralId")
                        if peripheral_id:
                            # 如果已有连接，先断开
                            if self.client and self.client.is_connected:
                                try:
                                    await self.client.disconnect()
                                except Exception as e:
                                    await log_exception_async(e, "connect断开旧连接")
                                    pass
                            
                            try:
                                self.client = BleakClient(peripheral_id, timeout=10.0)  # 添加超时
                                self.client.set_disconnected_callback(self.on_disconnect)
                                await self.client.connect()
                                
                                if self.client.is_connected:
                                    response = json.dumps({
                                        "jsonrpc": "2.0",
                                        "result": None,
                                        "id": request_id
                                    })
                                    await log_message("下发", response)
                                    if not websocket.closed:
                                        await websocket.send(response)
                                else:
                                    error_response = json.dumps({
                                        "jsonrpc": "2.0",
                                        "result": None,
                                        "id": request_id if request_id else 0
                                    })
                                    await log_message("下发", error_response)
                                    if not websocket.closed:
                                        await websocket.send(error_response)
                            except Exception as e:
                                await log_exception_async(e, "connect连接BLE")
                                error_response = json.dumps({
                                    "jsonrpc": "2.0",
                                    "result": None,
                                    "id": request_id if request_id else 0
                                })
                                await log_message("下发", error_response)
                                if not websocket.closed:
                                    await websocket.send(error_response)
                                if self.client:
                                    self.client = None
                                pass

                    elif method == "write":
                        service_id = params.get("serviceId")
                        characteristic_id = params.get("characteristicId")
                        message = params.get("message")
                        encoding = params.get("encoding", "utf-8")

                        if all([service_id, characteristic_id, message]) and self.client and self.client.is_connected:
                            try:
                                if encoding == "base64":
                                    message_bytes = base64.b64decode(message)
                                    print("write message_bytes",message_bytes)
                                else:
                                    print("write message",message)
                                    message_bytes = message.encode(encoding)

                                await self.client.write_gatt_char(characteristic_id, message_bytes)
                                response = json.dumps({
                                    "jsonrpc": "2.0",
                                    "result": None,
                                    "id": request_id
                                })
                                await log_message("下发", response)
                                if not websocket.closed:
                                    await websocket.send(response)
                            except Exception as e:
                                await log_exception_async(e, "write写入特征值")
                                error_response = json.dumps({
                                    "jsonrpc": "2.0",
                                    "result": None,
                                    "id": request_id if request_id else 0
                                })
                                await log_message("下发", error_response)
                                if not websocket.closed:
                                    await websocket.send(error_response)
                                pass
                        else:
                            error_response = json.dumps({
                                "jsonrpc": "2.0",
                                "result": None,
                                "id": request_id if request_id else 0
                            })
                            await log_message("下发", error_response)
                            if not websocket.closed:
                                await websocket.send(error_response)

                    elif method == "read":
                        service_id = params.get("serviceId")
                        characteristic_id = params.get("characteristicId")

                        if all([service_id, characteristic_id]) and self.client and self.client.is_connected:
                            try:
                                # 添加读取超时，避免长时间阻塞
                                data = await asyncio.wait_for(
                                    self.client.read_gatt_char(characteristic_id),
                                    timeout=10.0
                                )
                                print('read-data',data)
                                response = json.dumps({
                                    "jsonrpc": "2.0",
                                    "result": {
                                        "serviceId": service_id,
                                        "characteristicId": characteristic_id,
                                        "message": base64.b64encode(data).decode("utf-8")
                                    },
                                    "id": request_id
                                })
                                await log_message("下发", response)
                                # 检查WebSocket状态后再发送
                                if not websocket.closed:
                                    await websocket.send(response)
                            except asyncio.TimeoutError as e:
                                await log_exception_async(e, "read读取超时")
                                error_response = json.dumps({
                                    "jsonrpc": "2.0",
                                    "result": None,
                                    "id": request_id if request_id else 0
                                })
                                await log_message("下发", error_response)
                                if not websocket.closed:
                                    await websocket.send(error_response)
                                pass
                            except Exception as e:
                                await log_exception_async(e, "read读取特征值")
                                error_response = json.dumps({
                                    "jsonrpc": "2.0",
                                    "result": None,
                                    "id": request_id if request_id else 0
                                })
                                await log_message("下发", error_response)
                                # 检查WebSocket状态后再发送
                                if not websocket.closed:
                                    await websocket.send(error_response)
                                pass
                        else:
                            error_response = json.dumps({
                                "jsonrpc": "2.0",
                                "result": None,
                                "id": request_id if request_id else 0
                            })
                            await log_message("下发", error_response)
                            if not websocket.closed:
                                await websocket.send(error_response)

                    elif method == "startNotifications":
                        service_id = params.get("serviceId")
                        characteristic_id = params.get("characteristicId")

                        if all([service_id, characteristic_id]) and self.client and self.client.is_connected:
                            try:
                                # 如果已经监听，先停止旧的
                                if characteristic_id in self._notification_callbacks:
                                    try:
                                        await self.client.stop_notify(characteristic_id)
                                    except Exception as e:
                                        await log_exception_async(e, "startNotifications停止旧通知")
                                        pass
                                
                                # 创建新的通知处理器
                                callback = self.notification_handler(websocket, service_id, characteristic_id)
                                await self.client.start_notify(characteristic_id, callback)
                                self._notification_callbacks[characteristic_id] = callback
                                
                                response = json.dumps({
                                    "jsonrpc": "2.0",
                                    "result": None,
                                    "id": request_id
                                })
                                await log_message("下发", response)
                                if not websocket.closed:
                                    await websocket.send(response)
                            except Exception as e:
                                await log_exception_async(e, "startNotifications启动通知")
                                error_response = json.dumps({
                                    "jsonrpc": "2.0",
                                    "result": None,
                                    "id": request_id if request_id else 0
                                })
                                await log_message("下发", error_response)
                                if not websocket.closed:
                                    await websocket.send(error_response)
                                pass
                        else:
                            error_response = json.dumps({
                                "jsonrpc": "2.0",
                                "result": None,
                                "id": request_id if request_id else 0
                            })
                            await log_message("下发", error_response)
                            if not websocket.closed:
                                await websocket.send(error_response)
                            
                    elif method == "ping":
                        # 处理ping请求，返回pong响应
                        response = json.dumps({
                            "jsonrpc": "2.0",
                            "result": {"pong": True, "timestamp": int(time.time())},
                            "id": request_id if request_id else 0
                        })
                        await log_message("下发", response)
                        if not websocket.closed:
                            await websocket.send(response)

                except json.JSONDecodeError as e:
                    await log_exception_async(e, "JSON解析")
                    error_msg = json.dumps({
                        "jsonrpc": "2.0",
                        "result": None,
                        "id": request.get("id") if request else 0
                    })
                    await log_message("下发", error_msg)
                    try:
                        if not websocket.closed:
                            await websocket.send(error_msg)
                    except Exception as e2:
                        await log_exception_async(e2, "JSON解析后发送响应")
                        pass
                    pass
                except Exception as e:
                    await log_exception_async(e, "处理请求")
                    error_msg = json.dumps({
                        "jsonrpc": "2.0",
                        "result": None,
                        "id": request.get("id") if request else 0
                    })
                    await log_message("下发", error_msg)
                    try:
                        if not websocket.closed:
                            await websocket.send(error_msg)
                    except Exception as e2:
                        await log_exception_async(e2, "处理请求后发送响应")
                        pass
                    pass

        except websockets.exceptions.ConnectionClosed as e:
            await log_exception_async(e, "WebSocket连接关闭")
            pass
        finally:
            # 清理BLE连接和通知
            self._shutdown = True
            if self.client and self.client.is_connected:
                try:
                    # 停止所有通知
                    for characteristic_id in list(self._notification_callbacks.keys()):
                        try:
                            await self.client.stop_notify(characteristic_id)
                        except Exception as e:
                            await log_exception_async(e, f"停止通知{characteristic_id}")
                            pass
                    
                    # 断开连接
                    await self.client.disconnect()
                except Exception as e:
                    await log_exception_async(e, "断开BLE连接")
                    pass
                finally:
                    self.client = None
                    self.target_device = None
                    self._notification_callbacks.clear()
                    self.notification_records.clear()

    def notification_handler(self, websocket, service_id, characteristic_id):
        async def callback(sender, data):
            try:
                # 检查websocket是否已关闭
                if self.websocket is None or self.websocket.closed or self._shutdown:
                    return
                
                current_time = time.time()  # 使用time.time()而不是event_loop.time()
                last_message, last_time = self.notification_records[characteristic_id]
                
                # 解码当前数据用于比较
                current_message = base64.b64encode(data).decode('utf-8')
                print('notification_handler current_message',data)
                
                # 过滤逻辑 - 只在0.5秒内且消息完全相同时跳过
                if current_message == last_message and (current_time - last_time) < 0.5:
                    return
                
                # 更新记录
                self.notification_records[characteristic_id] = (current_message, current_time)
                
                response = json.dumps({
                    "jsonrpc": "2.0",
                    "method": "characteristicDidChange",
                    "params": {
                        "serviceId": service_id,
                        "characteristicId": characteristic_id,
                        "message": current_message
                    }
                })
                
                # 使用非阻塞日志
                try:
                    await log_message("下发", response)
                except Exception as e:
                    await log_exception_async(e, "notification_handler日志写入")
                    pass
                
                # 发送响应
                try:
                    await websocket.send(response)
                except websockets.exceptions.ConnectionClosed as e:
                    await log_exception_async(e, "notification_handler发送通知连接关闭")
                    self._shutdown = True
                    pass
                except Exception as e:
                    await log_exception_async(e, "notification_handler发送通知")
                    pass
            except Exception as e:
                await log_exception_async(e, "notification_handler")
                pass
        return callback

async def check_port_and_start_server(port=20111, host='localhost'):
    """检查端口并启动服务器"""
    if is_port_in_use(port, host):
        print(f"错误: 端口 {port} 已被占用，无法启动服务")
        return False
    
    print(f"端口 {port} 可用，正在启动服务...")
    server = await websockets.serve(
        lambda websocket, path: BLEClient().handle_client(websocket, path),
        host, port,
        ping_interval=20,  # 每20秒发送ping
        ping_timeout=10,   # 等待pong超时10秒
        max_size=2 * 1024 * 1024,  # 最大消息大小2MB
        close_timeout=5,   # 关闭超时5秒
    )
    
    print(f"WebSocket服务已启动: ws://{host}:{port}/scratch/ble")
    print("日志文件路径: ./b.log")
    
    # 执行自检测试
    try:
        async with websockets.connect(f"ws://{host}:{port}/scratch/ble") as websocket:
            print("正在执行自检测试...")
            test_result = await self_test(websocket)
            if test_result:
                print("自检测试成功: 服务正常运行")
            else:
                print("自检测试失败: 服务可能存在问题")
    except Exception as e:
        await log_exception_async(e, "自检测试")
        pass
    
    return server

async def main():
    server = await check_port_and_start_server()
    if server:
        try:
            await asyncio.Future()  # 保持服务器运行
        finally:
            # 清理日志系统
            global _log_queue, _log_task
            if _log_queue is not None:
                try:
                    await _log_queue.put(None)
                except Exception as e:
                    await log_exception_async(e, "main清理日志队列")
                    pass
            if _log_task is not None:
                try:
                    await _log_task
                except Exception as e:
                    await log_exception_async(e, "main等待日志任务")
                    pass

if __name__ == "__main__":
    asyncio.run(main())