# -*- 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('a.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:
            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} 条日志")

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:
        await log_message("自测", "自测超时，未收到响应")
        return False
    except Exception as e:
        await log_message("自测", f"自测异常: {str(e)}")
        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  # 当前扫描器实例
        self._device_name_from_advertisement = None  # 从广播包响应中获取的设备名称

    def on_disconnect(self, client):
        print("BLE连接断开，关闭WebSocket")
        # 异步记录日志
        asyncio.create_task(log_message("系统", "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:
                # 如果没有事件循环，直接设置为None（延迟清理）
                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连接已关闭")
                await log_message("系统", "WebSocket连接已关闭")
            except Exception as e:
                print(f"关闭WebSocket时出错: {e}")
                await log_message("错误", f"关闭WebSocket时出错: {e}")
            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更准确
            device_name_from_adv = advertisement_data.local_name if advertisement_data.local_name else None
            
            # 如果广播包中没有名称，再使用device.name作为备用
            device_name = device_name_from_adv if device_name_from_adv else (device.name if device.name else "")
            
            # 丢弃名字为空的设备
            if not device_name or device_name.strip() == "":
                print(f"跳过名字为空的设备: {device.address}")
                # 异步记录日志
                asyncio.create_task(log_message("扫描", f"跳过名字为空的设备: {device.address}"))
                return None
            
            # 保存从广播包响应中获取的设备名称
            self._device_name_from_advertisement = device_name_from_adv if device_name_from_adv else 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_from_adv if device_name_from_adv else 'N/A'}")
                print(f"设备名称(备用来源): {device.name if device.name else 'N/A'}")
                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'}")
                # 异步记录设备发现日志
                device_info = f"找到目标设备: {device_name} ({device.address}), RSSI: {device.rssi} dBm"
                asyncio.create_task(log_message("扫描", device_info))
                # 发现设备后，尝试立即停止扫描
                try:
                    if self._scanner is not None:
                        # 在事件循环中停止扫描，避免阻塞回调
                        asyncio.create_task(self._stop_scan_early())
                except Exception:
                    pass
            return self.target_device

    async def _stop_scan_early(self):
        try:
            if self._scanner is not None:
                await self._scanner.stop()
        except Exception:
            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", [])
                        self._device_name_from_advertisement = None  # 重置保存的设备名称

                        # 双重扫描：快速扫描后若未发现，再进行扩展扫描；发现即停
                        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）...")
                            await log_message("扫描", 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:
                                    pass
                                self._scanner = None
                            
                            if self.target_device:
                                found = True
                                break
                            else:
                                print(f"第{phase_index}阶段未找到设备")
                                await log_message("扫描", 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 None
                            
                            # 如果广播包响应中没有名称，使用之前保存的名称或device.name作为备用
                            if not device_name:
                                device_name = self._device_name_from_advertisement
                            if not device_name:
                                device_name = device.name if device.name else None
                            
                            # 补充检查：如果设备名称仍为空，使用设备地址作为名称
                            if not device_name or device_name.strip() == "":
                                print(f"警告: 设备名称为空，使用设备地址作为名称: {device.address}")
                                await log_message("警告", f"设备名称为空，使用设备地址作为名称: {device.address}")
                                device_name = device.address  # 使用设备地址作为备用名称
                            
                            # 记录名称来源
                            name_source = "广播包响应(local_name)" if adv_data.local_name else \
                                         ("已保存的广播包名称" if self._device_name_from_advertisement else \
                                         ("device.name" if device.name else "设备地址"))
                            print(f"响应包中使用的设备名称来源: {name_source}")
                            await log_message("系统", f"响应包中使用的设备名称来源: {name_source}")
                            
                            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:
                                    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:
                                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

                    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)
                                    await log_message("写入", f"Base64编码数据: {message_bytes.hex() if len(message_bytes) < 100 else message_bytes[:50].hex() + '...'}")
                                else:
                                    print("write message",message)
                                    message_bytes = message.encode(encoding)
                                    await log_message("写入", f"UTF-8消息: {message}")

                                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:
                                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)
                        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)
                                await log_message("读取", f"特征 {characteristic_id}: {data.hex() if len(data) < 100 else data[:50].hex() + '...'}")
                                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:
                                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:
                                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)
                        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:
                                        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:
                                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)
                        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:
                    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:
                        pass
                except Exception as e:
                    # 记录完整错误堆栈
                    error_trace = traceback.format_exc()
                    print(f"处理请求时出错: {error_trace}")
                    await log_message("错误", f"处理请求时出错: {error_trace}")
                    
                    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:
                        pass

        except websockets.exceptions.ConnectionClosed:
            print("WebSocket连接关闭")
            await log_message("系统", "WebSocket连接关闭")
        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:
                            print(f"停止通知失败 {characteristic_id}: {e}")
                            await log_message("错误", f"停止通知失败 {characteristic_id}: {e}")
                    
                    # 断开连接
                    await self.client.disconnect()
                except Exception as e:
                    print(f"断开BLE连接失败: {e}")
                    await log_message("错误", f"断开BLE连接失败: {e}")
                finally:
                    self.client = None
                    self.target_device = None
                    self._device_name_from_advertisement = 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)
                # 异步记录通知日志（不阻塞通知处理）
                asyncio.create_task(log_message("通知", f"特征 {characteristic_id}: {data.hex() if len(data) < 100 else data[:50].hex() + '...'}"))
                
                # 过滤逻辑 - 只在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:
                    print(f"日志写入失败: {e}")
                
                # 发送响应
                try:
                    await websocket.send(response)
                except websockets.exceptions.ConnectionClosed:
                    print("WebSocket连接已关闭，停止发送通知")
                    await log_message("系统", "WebSocket连接已关闭，停止发送通知")
                    self._shutdown = True
                except Exception as e:
                    print(f"发送通知失败: {e}")
                    await log_message("错误", f"发送通知失败: {e}")
            except Exception as e:
                print(f"通知处理器出错: {e}")
                # 注意：这里不能使用 await，因为可能不在 async 上下文中
                try:
                    asyncio.create_task(log_message("错误", f"通知处理器出错: {e}"))
                except Exception:
                    pass
        return callback

async def check_port_and_start_server(port=20111, host='localhost'):
    """检查端口并启动服务器"""
    if is_port_in_use(port, host):
        print(f"错误: 端口 {port} 已被占用，无法启动服务")
        await log_message("错误", f"端口 {port} 已被占用，无法启动服务")
        return False
    
    print(f"端口 {port} 可用，正在启动服务...")
    await log_message("系统", 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("日志文件路径: ./a.log")
    await log_message("系统", f"WebSocket服务已启动: ws://{host}:{port}/scratch/ble")
    
    # 执行自检测试
    try:
        async with websockets.connect(f"ws://{host}:{port}/scratch/ble") as websocket:
            print("正在执行自检测试...")
            await log_message("系统", "正在执行自检测试...")
            test_result = await self_test(websocket)
            if test_result:
                print("自检测试成功: 服务正常运行")
                await log_message("系统", "自检测试成功: 服务正常运行")
            else:
                print("自检测试失败: 服务可能存在问题")
                await log_message("警告", "自检测试失败: 服务可能存在问题")
    except Exception as e:
        print(f"自检测试异常: {str(e)}")
        await log_message("错误", f"自检测试异常: {str(e)}")
    
    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:
                    pass
            if _log_task is not None:
                try:
                    await _log_task
                except Exception:
                    pass

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