# server.py
import asyncio
import websockets
from datetime import datetime
from urllib.parse import urlparse, parse_qs

clients = set()
mac_to_websocket = {}
mac_lock = asyncio.Lock()


async def handle_client(websocket):
    clients.add(websocket)
    client_ip = websocket.remote_address[0]

    # 解析MAC地址
    parsed_path = urlparse(websocket.path)
    query_params = parse_qs(parsed_path.query)
    mac_addresses = query_params.get("mac", [])
    if not mac_addresses:
        print(f"[{datetime.now().strftime('%H:%M:%S')}] 客户端未提供MAC地址，断开连接")
        clients.remove(websocket)
        await websocket.close()
        return
    mac_address = mac_addresses[0]

    # 处理已存在的同MAC连接
    async with mac_lock:
        if mac_address in mac_to_websocket:
            old_ws = mac_to_websocket[mac_address]
            if old_ws.open:
                await old_ws.close()
            clients.remove(old_ws)
        mac_to_websocket[mac_address] = websocket

    print(
        f"[{datetime.now().strftime('%H:%M:%S')}] 新客户端接入 | MAC: {mac_address} | IP: {client_ip} | 总数: {len(clients)}"
    )

    try:
        async for message in websocket:
            print(
                f"[{datetime.now().strftime('%H:%M:%S')}] 收到来自 {mac_address} 的消息: {message}"
            )
            # 解析目标MAC地址和消息内容
            if "|" not in message:
                print(f"无效消息格式: {message}")
                continue
            target_mac, content = message.split("|", 1)
            async with mac_lock:
                target_ws = mac_to_websocket.get(target_mac)
            if target_ws and target_ws.open:
                await target_ws.send(f"[{mac_address}] {content}")
            else:
                print(f"目标MAC {target_mac} 未找到或连接已关闭")
    finally:
        async with mac_lock:
            if (
                mac_address in mac_to_websocket
                and mac_to_websocket[mac_address] == websocket
            ):
                del mac_to_websocket[mac_address]
        clients.remove(websocket)
        print(
            f"[{datetime.now().strftime('%H:%M:%S')}] 客户端断开 | MAC: {mac_address} | 剩余: {len(clients)}"
        )


async def main():
    async with websockets.serve(handle_client, "0.0.0.0", 8765):
        print("WebSocket服务器已启动 ws://0.0.0.0:8765")
        await asyncio.Future()  # 永久运行


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