import asyncio
import json
import logging
import signal
from typing import Any, TypeVar

import websockets
from pydantic import BaseModel
from websockets.server import WebSocketServerProtocol, serve

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

R = TypeVar("R", bound=BaseModel)


class Event(BaseModel):
    """强制使用该结构传递数据"""

    type: str
    data: str | dict[str, Any]


class TokenEvent(Event):
    type: str = "token"
    data: str


USERS = {"123": "xiao"}


def get_user(token: str):
    """Find user authenticated by token or return None."""
    return USERS.get(token)


async def parse_event(websocket: WebSocketServerProtocol, model: type[R]) -> R | None:
    data = await websocket.recv()
    try:
        return model.parse_raw(data)
    except Exception as e:
        logging.error(e)
        await websocket.send(json.dumps({"type": "error", "data": "parameter error"}))


async def handler(websocket: WebSocketServerProtocol) -> None:
    data = await parse_event(websocket, TokenEvent)
    if data is None or data.type != "token":
        # 首次进入必须传递 token
        return
    # 判断 user
    user = get_user(data.data)
    if user is None:
        await websocket.close(1011, "authentication failed")
        return
    print(user)
    async for message in websocket:
        # 之后不再验证 token
        try:
            data = Event.parse_raw(message)
        except Exception as e:
            logging.error(e)
            # 客户端需要处理错误后续操作，关闭连接还是重试
            await websocket.send(json.dumps({"type": "error", "data": "参数错误"}))


async def main(port: int = 8765) -> None:
    # set the stop condition when receiving SIGINT or SIGTERM.
    loop = asyncio.get_running_loop()
    stop = loop.create_future()
    loop.add_signal_handler(signal.SIGINT, stop.set_result, None)
    loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)

    async with serve(handler, "localhost", port):
        logging.info(f"Running server on port {port}")
        await stop
        logging.info("exit")


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

async with serve(
    handler,
    "localhost",
    port,
    ping_interval=20,  # 20 秒 ping 一次， None 禁用
    ping_timeout=20,  # 20 秒内没有收到 pong，关闭连接(1011)
)

CLIENTS = set()

async def handler(websocket: WebSocketServerProtocol) -> None:
    # 加入客户端列表
    CLIENTS.add(websocket)
    # 向所有客户端发送消息
    websockets.broadcast(CLIENTS, f"clients num: {len(CLIENTS)}")
    try:
        async for message in websocket:
            print(message)
    finally:
        CLIENTS.remove(websocket)
        # 向所有客户端再次发送消息
   	    websockets.broadcast(CLIENTS, f"clients num: {len(CLIENTS)}")

websockets.serve(
    ...,
    compression=None,  # None 禁用
)

# 自定义压缩设置
from websockets.extensions import permessage_deflate
websockets.serve(
    ...,
    extensions=[
        permessage_deflate.ServerPerMessageDeflateFactory(
            server_max_window_bits=12,  # 9-15之间
            client_max_window_bits=12,
            compress_settings={"memLevel": 5},
        ),
    ],
)

async with websockets.unix_serve(
    echo,
    path="websocket.sock",
):
    await stop