fin_data_unify_plat/backend/app/api/v1/realtime.py

"""
实时行情 API 路由
"""
import json
import logging
from typing import Annotated, List, Optional
from datetime import datetime

from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, Query, HTTPException
from sqlalchemy.orm import Session

from app.schemas import (
    RealtimeQuoteItem,
    SubscribeRequest,
    UnsubscribeRequest,
    ResponseData
)
from app.services.realtime_service import realtime_service
from app.services.auth_service import decode_token
from app.db.init_db import get_sqlite_db

router = APIRouter()
logger = logging.getLogger(__name__)

# WebSocket 连接数限制配置
MAX_CONNECTIONS_PER_USER = 5
MAX_TOTAL_CONNECTIONS = 100


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """
    WebSocket 连接 - 实时行情推送
    
    连接后发送订阅消息:
    {"action": "subscribe", "symbols": ["IF2406", "IC2406"]}
    
    取消订阅:
    {"action": "unsubscribe", "symbols": ["IF2406"]}
    
    心跳:
    {"action": "ping"}
    """
    await websocket.accept()
    subscribed_symbols = set()
    user_id: Optional[int] = None
    token: Optional[str] = None
    
    try:
        # 尝试从连接参数或第一个消息中获取 Token 进行认证
        # 先等待认证消息
        data = await websocket.receive_text()
        message = json.loads(data)
        
        # 检查是否是认证消息
        if message.get("action") == "auth":
            token = message.get("token")
            if token:
                try:
                    payload = decode_token(token)
                    user_id = payload.get("sub")
                    logger.info(f"WebSocket user {user_id} authenticated")
                except Exception as e:
                    logger.warning(f"WebSocket auth failed: {e}")
                    await websocket.send_text(json.dumps({
                        "type": "error",
                        "message": "认证失败"
                    }))
                    await websocket.close(code=4001, reason="认证失败")
                    return
        else:
            # 如果没有认证，使用匿名连接（限制更严格）
            logger.info("WebSocket connected without authentication")
        
        # 检查总连接数限制
        total_connections = realtime_service.get_total_connections()
        if total_connections >= MAX_TOTAL_CONNECTIONS:
            await websocket.send_text(json.dumps({
                "type": "error",
                "message": "服务器连接数已达上限"
            }))
            await websocket.close(code=4002, reason="连接数超限")
            return
        
        # 检查单用户连接数限制
        if user_id:
            user_connections = realtime_service.get_user_connections(user_id)
            if len(user_connections) >= MAX_CONNECTIONS_PER_USER:
                await websocket.send_text(json.dumps({
                    "type": "error",
                    "message": f"单用户最大连接数限制 ({MAX_CONNECTIONS_PER_USER})"
                }))
                await websocket.close(code=4003, reason="连接数超限")
                return
        
        # 注册连接
        realtime_service.register_connection(websocket, user_id)
        
        # 处理第一条消息（如果不是认证消息）
        if message.get("action") != "auth":
            action = message.get("action")
            symbols = message.get("symbols", [])
            
            if action == "subscribe":
                for symbol in symbols:
                    await realtime_service.subscribe_symbol(symbol, websocket)
                    subscribed_symbols.add(symbol)
                
                await websocket.send_text(json.dumps({
                    "type": "subscribed",
                    "symbols": list(subscribed_symbols)
                }))
                
                for symbol in symbols:
                    quote = await realtime_service.get_latest_quote(symbol)
                    if quote:
                        await websocket.send_text(json.dumps({
                            "type": "quote",
                            "symbol": symbol,
                            "data": quote
                        }))
        
        # 主消息循环
        while True:
            data = await websocket.receive_text()
            message = json.loads(data)
            
            action = message.get("action")
            symbols = message.get("symbols", [])
            
            if action == "subscribe":
                for symbol in symbols:
                    await realtime_service.subscribe_symbol(symbol, websocket)
                    subscribed_symbols.add(symbol)
                
                await websocket.send_text(json.dumps({
                    "type": "subscribed",
                    "symbols": list(subscribed_symbols)
                }))
                
                for symbol in symbols:
                    quote = await realtime_service.get_latest_quote(symbol)
                    if quote:
                        await websocket.send_text(json.dumps({
                            "type": "quote",
                            "symbol": symbol,
                            "data": quote
                        }))
            
            elif action == "unsubscribe":
                for symbol in symbols:
                    await realtime_service.unsubscribe_symbol(symbol, websocket)
                    subscribed_symbols.discard(symbol)
                
                await websocket.send_text(json.dumps({
                    "type": "unsubscribed",
                    "symbols": symbols
                }))
            
            elif action == "ping":
                await websocket.send_text(json.dumps({
                    "type": "pong",
                    "timestamp": datetime.utcnow().isoformat()
                }))
            
            elif action == "auth":
                # 重新认证
                token = message.get("token")
                if token:
                    try:
                        payload = decode_token(token)
                        user_id = payload.get("sub")
                        logger.info(f"WebSocket user {user_id} re-authenticated")
                        await websocket.send_text(json.dumps({
                            "type": "authenticated",
                            "user_id": user_id
                        }))
                    except Exception as e:
                        await websocket.send_text(json.dumps({
                            "type": "error",
                            "message": f"认证失败：{str(e)}"
                        }))
            
            else:
                await websocket.send_text(json.dumps({
                    "type": "error",
                    "message": f"Unknown action: {action}"
                }))
    
    except WebSocketDisconnect:
        logger.info(f"WebSocket client disconnected (user={user_id})")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
    finally:
        # 清理订阅和连接
        for symbol in subscribed_symbols:
            await realtime_service.unsubscribe_symbol(symbol, websocket)
        if user_id:
            realtime_service.unregister_connection(websocket, user_id)
        else:
            realtime_service.unregister_anonymous_connection(websocket)


@router.get("/quote", response_model=ResponseData)
async def get_latest_quote(
    symbol: Annotated[str, Query(description="品种代码")]
):
    """获取最新行情"""
    quote = await realtime_service.get_latest_quote(symbol)
    if not quote:
        return ResponseData(
            code=404,
            message="Quote not found",
            data=None
        )
    
    return ResponseData(
        code=0,
        message="success",
        data=quote
    )


@router.get("/quotes", response_model=ResponseData)
async def get_multiple_quotes(
    symbols: Annotated[str, Query(description="品种代码列表，逗号分隔")]
):
    """获取多个品种的最新行情"""
    symbol_list = [s.strip() for s in symbols.split(",")]
    quotes = {}
    
    for symbol in symbol_list:
        quote = await realtime_service.get_latest_quote(symbol)
        if quote:
            quotes[symbol] = quote
    
    return ResponseData(
        code=0,
        message="success",
        data={"quotes": quotes, "count": len(quotes)}
    )


@router.get("/subscriptions", response_model=ResponseData)
async def get_active_subscriptions():
    """获取活跃订阅统计"""
    stats = realtime_service.get_active_subscriptions()
    return ResponseData(
        code=0,
        message="success",
        data={
            "subscriptions": stats,
            "total_symbols": len(stats),
            "total_connections": sum(stats.values())
        }
    )