sanguo_vnpy_v2/sanguo_web/websocket/routes.py

"""
WebSocket 路由
处理 WebSocket 连接和实时数据推送
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from typing import Optional
import logging
from datetime import datetime

from .manager import manager

logger = logging.getLogger(__name__)

router = APIRouter()


@router.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    token: Optional[str] = Query(None, description="JWT Token")
):
    """
    WebSocket 端点

    - **token**: 可选的 JWT Token，用于身份验证

    支持的消息类型:
    - subscribe: 订阅数据类型（tick, order, trade, position, account, log, contract）
    - unsubscribe: 取消订阅
    - subscribe_symbol: 订阅特定品种的行情数据
    - unsubscribe_symbol: 取消品种订阅
    - ping: 心跳检测
    """
    # 验证 Token（可选）
    user = "anonymous"
    if token:
        try:
            from ..deps import verify_token
            payload = verify_token(token)
            user = payload.get("sub", "anonymous")
        except Exception as e:
            logger.warning(f"WebSocket authentication failed: {e}")
            await websocket.close(code=1008, reason="Invalid token")
            return

    # 接受连接
    connection_id = await manager.connect(websocket, user)

    try:
        # 发送欢迎消息
        await manager.send_personal_message({
            "type": "connected",
            "connection_id": connection_id,
            "user": user,
            "timestamp": datetime.now().isoformat()
        }, connection_id)

        # 处理消息循环
        while True:
            data = await websocket.receive_json()

            message_type = data.get("type")
            message_data = data.get("data", {})

            if message_type == "subscribe":
                # 订阅数据类型
                subscription_type = message_data.get("subscription")
                if subscription_type:
                    # 支持单个订阅或列表订阅
                    if isinstance(subscription_type, list):
                        subscribed = []
                        for sub_type in subscription_type:
                            if manager.subscribe(connection_id, sub_type):
                                subscribed.append(sub_type)
                        await manager.send_personal_message({
                            "type": "subscribed",
                            "subscriptions": subscribed
                        }, connection_id)
                    else:
                        if manager.subscribe(connection_id, subscription_type):
                            await manager.send_personal_message({
                                "type": "subscribed",
                                "subscription": subscription_type
                            }, connection_id)

            elif message_type == "unsubscribe":
                # 取消订阅
                subscription_type = message_data.get("subscription")
                if subscription_type:
                    # 支持单个取消或列表取消
                    if isinstance(subscription_type, list):
                        unsubscribed = []
                        for sub_type in subscription_type:
                            if manager.unsubscribe(connection_id, sub_type):
                                unsubscribed.append(sub_type)
                        await manager.send_personal_message({
                            "type": "unsubscribed",
                            "subscriptions": unsubscribed
                        }, connection_id)
                    else:
                        if manager.unsubscribe(connection_id, subscription_type):
                            await manager.send_personal_message({
                                "type": "unsubscribed",
                                "subscription": subscription_type
                            }, connection_id)

            elif message_type == "subscribe_symbol":
                # 订阅品种
                symbol = message_data.get("symbol")
                if symbol:
                    # 支持单个品种或列表订阅
                    if isinstance(symbol, list):
                        subscribed_symbols = []
                        for sym in symbol:
                            if manager.subscribe_symbol(connection_id, sym):
                                subscribed_symbols.append(sym)
                        await manager.send_personal_message({
                            "type": "symbol_subscribed",
                            "symbols": subscribed_symbols
                        }, connection_id)
                    else:
                        if manager.subscribe_symbol(connection_id, symbol):
                            await manager.send_personal_message({
                                "type": "symbol_subscribed",
                                "symbol": symbol
                            }, connection_id)

            elif message_type == "unsubscribe_symbol":
                # 取消品种订阅
                symbol = message_data.get("symbol")
                if symbol:
                    # 支持单个品种或列表取消
                    if isinstance(symbol, list):
                        unsubscribed_symbols = []
                        for sym in symbol:
                            if manager.unsubscribe_symbol(connection_id, sym):
                                unsubscribed_symbols.append(sym)
                        await manager.send_personal_message({
                            "type": "symbol_unsubscribed",
                            "symbols": unsubscribed_symbols
                        }, connection_id)
                    else:
                        if manager.unsubscribe_symbol(connection_id, symbol):
                            await manager.send_personal_message({
                                "type": "symbol_unsubscribed",
                                "symbol": symbol
                            }, connection_id)

            elif message_type == "ping":
                # 心跳
                await manager.send_personal_message({
                    "type": "pong"
                }, connection_id)

            else:
                await manager.send_personal_message({
                    "type": "error",
                    "message": f"Unknown message type: {message_type}"
                }, connection_id)

    except WebSocketDisconnect:
        logger.info(f"WebSocket disconnected normally: {connection_id}")
    except Exception as e:
        logger.error(f"WebSocket error for {connection_id}: {e}")
    finally:
        await manager.disconnect(connection_id)


@router.get("/ws/status")
async def websocket_status():
    """
    获取 WebSocket 连接状态
    """
    return {
        "active_connections": manager.get_connection_count(),
        "users": len(manager.user_connections)
    }