"""
实时数据路由
"""
from fastapi import APIRouter, Depends, Query, WebSocket, WebSocketDisconnect
from sqlalchemy.orm import Session

from app.db.session import get_db
from app.schemas.base import ResponseModel
from app.services.realtime_service import RealtimeService
from app.core.security import get_current_user
from app.models.user import User

router = APIRouter()


@router.get("/snapshot", response_model=ResponseModel)
async def get_realtime_snapshot(
    codes: str = Query(..., description="代码列表，多个用逗号分隔"),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """获取最新快照数据"""
    service = RealtimeService(db)
    code_list = [c.strip() for c in codes.split(",")]
    data = service.get_latest_snapshot(code_list)
    return ResponseModel(data=data)


@router.post("/subscribe", response_model=ResponseModel)
async def subscribe_realtime(
    codes: str = Query(..., description="代码列表，多个用逗号分隔"),
    types: str = Query("snapshot", description="订阅类型: snapshot, kline"),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """开始实时数据订阅"""
    code_list = [c.strip() for c in codes.split(",")]
    return ResponseModel(data={
        "message": "订阅成功",
        "codes": code_list,
        "types": types.split(",")
    })


@router.delete("/subscribe", response_model=ResponseModel)
async def unsubscribe_realtime(
    codes: str = Query(None, description="代码列表，多个用逗号分隔，为空则取消所有"),
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user)
):
    """停止实时数据订阅"""
    code_list = [c.strip() for c in codes.split(",")] if codes else None
    return ResponseModel(data={
        "message": "取消订阅成功",
        "codes": code_list
    })


@router.get("/subscribe/status", response_model=ResponseModel)
async def get_subscribe_status(
    current_user: User = Depends(get_current_user)
):
    """获取订阅状态"""
    return ResponseModel(data={
        "subscribed_codes": [],
        "subscribed_types": []
    })


@router.websocket("/stream")
async def realtime_websocket(websocket: WebSocket):
    """WebSocket实时数据流"""
    # 获取查询参数
    query_params = websocket.query_params
    codes_str = query_params.get("codes", "")
    types_str = query_params.get("types", "snapshot")
    
    codes = [c.strip() for c in codes_str.split(",") if c.strip()]
    
    if not codes:
        await websocket.close(code=4000, reason="Missing codes parameter")
        return
    
    # 创建数据库会话
    from app.db.session import SessionLocal
    db = SessionLocal()
    
    try:
        service = RealtimeService(db)
        await service.subscribe_websocket(websocket, codes)
        
        # 保持连接
        while True:
            try:
                # 接收客户端消息（心跳检测）
                data = await websocket.receive_text()
                # 回复心跳
                await websocket.send_json({"type": "heartbeat", "timestamp": datetime.utcnow().isoformat()})
            except WebSocketDisconnect:
                break
            except Exception as e:
                break
    
    finally:
        await service.unsubscribe_websocket(websocket, codes)
        db.close()


from datetime import datetime