baoxiang/backend/app/routers/websocket.py

"""
WebSocket路由 - 实时通信
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
from typing import Dict, List
import json
import asyncio
from datetime import datetime
from ..core.security import verify_token
from ..utils.redis import get_pool_cache
from ..services.game_service import GameService
from ..services.countdown_service import countdown_manager
from ..core.database import SessionLocal

router = APIRouter(tags=["websocket"])

# 连接管理器
class ConnectionManager:
    def __init__(self):
        # 主播ID -> 连接列表
        self.streamer_connections: Dict[int, List[WebSocket]] = {}
        # 用户ID -> 连接列表
        self.user_connections: Dict[int, List[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, streamer_id: int = None, user_id: int = None):
        await websocket.accept()

        if streamer_id:
            if streamer_id not in self.streamer_connections:
                self.streamer_connections[streamer_id] = []
            self.streamer_connections[streamer_id].append(websocket)

        if user_id:
            if user_id not in self.user_connections:
                self.user_connections[user_id] = []
            self.user_connections[user_id].append(websocket)

    def disconnect(self, websocket: WebSocket, streamer_id: int = None, user_id: int = None):
        if streamer_id and streamer_id in self.streamer_connections:
            if websocket in self.streamer_connections[streamer_id]:
                self.streamer_connections[streamer_id].remove(websocket)
        if user_id and user_id in self.user_connections:
            if websocket in self.user_connections[user_id]:
                self.user_connections[user_id].remove(websocket)

    async def send_to_streamer(self, streamer_id: int, message: dict):
        if streamer_id in self.streamer_connections:
            for connection in self.streamer_connections[streamer_id]:
                try:
                    await connection.send_text(json.dumps(message, ensure_ascii=False))
                except:
                    pass

    async def send_to_user(self, user_id: int, message: dict):
        if user_id in self.user_connections:
            for connection in self.user_connections[user_id]:
                try:
                    await connection.send_text(json.dumps(message, ensure_ascii=False))
                except:
                    pass

    async def broadcast_to_all(self, message: dict):
        """广播给所有连接"""
        all_connections = []
        for connections in self.streamer_connections.values():
            all_connections.extend(connections)
        for connections in self.user_connections.values():
            all_connections.extend(connections)

        for connection in all_connections:
            try:
                await connection.send_text(json.dumps(message, ensure_ascii=False))
            except:
                pass


manager = ConnectionManager()


@router.websocket("/ws/streamer/{streamer_id}")
async def websocket_endpoint_for_streamer(
    websocket: WebSocket,
    streamer_id: int,
    token: str = Query(...)
):
    # 验证令牌
    user_id = verify_token(token)
    if not user_id:
        await websocket.close(code=4001)
        return

    await manager.connect(websocket, streamer_id=streamer_id)

    print(f"Streamer {streamer_id} WebSocket connected")

    try:
        # 启动奖池广播任务（频率降低到5秒）
        broadcast_task = asyncio.create_task(
            broadcast_pool_updates(streamer_id)
        )
        print(f"Started pool broadcast task for streamer {streamer_id}")

        while True:
            # 保持连接
            data = await websocket.receive_text()
            # 可以处理来自客户端的消息

    except WebSocketDisconnect:
        manager.disconnect(websocket, streamer_id=streamer_id)
        if 'broadcast_task' in locals():
            broadcast_task.cancel()
    except Exception as e:
        print(f"WebSocket error: {e}")
        manager.disconnect(websocket, streamer_id=streamer_id)
        if 'broadcast_task' in locals():
            broadcast_task.cancel()


@router.websocket("/ws/user/{user_id}")
async def websocket_endpoint_for_user(
    websocket: WebSocket,
    user_id: int,
    token: str = Query(...)
):
    # 验证令牌
    token_user_id = verify_token(token)
    if not token_user_id or int(token_user_id) != user_id:
        await websocket.close(code=4001)
        return

    await manager.connect(websocket, user_id=user_id)

    try:
        while True:
            data = await websocket.receive_text()
            # 处理用户消息

    except WebSocketDisconnect:
        manager.disconnect(websocket, user_id=user_id)
    except Exception as e:
        print(f"WebSocket error: {e}")
        manager.disconnect(websocket, user_id=user_id)


# 添加通用WebSocket端点来处理查询参数
@router.websocket("/socket.io/")
async def websocket_endpoint(
    websocket: WebSocket,
    role: str = Query(...),
    id: int = Query(...),
    token: str = Query(...)
):
    # 验证令牌
    user_id = verify_token(token)
    if not user_id:
        await websocket.close(code=4001)
        return

    # 根据角色连接到相应的管理器
    if role == "streamer":
        streamer_id = int(id)
        await manager.connect(websocket, streamer_id=streamer_id)

        print(f"Streamer {streamer_id} WebSocket (socket.io) connected")

        try:
            # 启动奖池广播任务（频率降低到5秒）
            broadcast_task = asyncio.create_task(
                broadcast_pool_updates(streamer_id)
            )
            print(f"Started pool broadcast task (socket.io) for streamer {streamer_id}")

            while True:
                data = await websocket.receive_text()
                # 可以处理来自客户端的消息

        except WebSocketDisconnect:
            manager.disconnect(websocket, streamer_id=streamer_id)
            if 'broadcast_task' in locals():
                broadcast_task.cancel()
        except Exception as e:
            print(f"WebSocket error: {e}")
            manager.disconnect(websocket, streamer_id=streamer_id)
            if 'broadcast_task' in locals():
                broadcast_task.cancel()

    elif role == "user":
        user_id = int(id)
        await manager.connect(websocket, user_id=user_id)

        try:
            while True:
                data = await websocket.receive_text()
                # 处理用户消息

        except WebSocketDisconnect:
            manager.disconnect(websocket, user_id=user_id)
        except Exception as e:
            print(f"WebSocket error: {e}")
            manager.disconnect(websocket, user_id=user_id)
    else:
        await websocket.close(code=4000)


async def broadcast_pool_updates(streamer_id: int):
    """定期广播奖池更新（不包含倒计时）"""
    print(f"Broadcast pool updates task started for streamer {streamer_id}")
    while True:
        try:
            # 获取主播的所有宝箱
            db = SessionLocal()
            chests = GameService.get_active_chests(db, streamer_id)

            for chest in chests:
                # 获取最新奖池数据
                pool_data = get_pool_cache(chest.id)

                # 计算赔率 - 修改为只对获胜方抽水
                total = pool_data["pool_a"] + pool_data["pool_b"]
                if total > 0:
                    # 赔率计算改为：(自己奖池 + 对方奖池 * 0.9) / 自己奖池
                    odds_a = round((pool_data["pool_a"] + pool_data["pool_b"] * 0.9) / pool_data["pool_a"], 2) if pool_data["pool_a"] > 0 else 0
                    odds_b = round((pool_data["pool_b"] + pool_data["pool_a"] * 0.9) / pool_data["pool_b"], 2) if pool_data["pool_b"] > 0 else 0
                else:
                    odds_a = odds_b = 0

                # 发送奖池更新消息
                pool_message = {
                    "type": "pool_update",
                    "chest_id": chest.id,
                    "pool_a": pool_data["pool_a"],
                    "pool_b": pool_data["pool_b"],
                    "odds_a": odds_a,
                    "odds_b": odds_b,
                    "total_bets": chest.total_bets
                }

                await manager.send_to_streamer(streamer_id, pool_message)

            db.close()

            # 每5秒广播一次奖池更新（大幅降低频率，节省资源）
            await asyncio.sleep(5)

        except Exception as e:
            print(f"Broadcast error: {e}")
            await asyncio.sleep(1)


# ==================== 倒计时相关函数（已废弃） ====================
# 说明：前端已实现本地倒计时，后端通过SchedulerService自动封盘
#       以下函数不再使用，保留以供参考

# async def start_chests_countdown(streamer_id: int):
#     """启动主播所有活跃宝箱的倒计时（已废弃）"""
#     print(f"Starting countdown for streamer {streamer_id}")
#     try:
#         db = SessionLocal()
#         chests = GameService.get_active_chests(db, streamer_id)
#
#         print(f"Found {len(chests)} chests for streamer {streamer_id}")
#
#         for chest in chests:
#             print(f"Checking chest {chest.id}, status: {chest.status}")
#             if chest.status == 0:  # BETTING
#                 print(f"Starting countdown for chest {chest.id}")
#                 # 启动倒计时
#                 await countdown_manager.start_chest_countdown(
#                     chest,
#                     on_update=lambda cid, time_rem: send_countdown_update(cid, time_rem, streamer_id),
#                     on_expire=lambda cid: handle_countdown_expire(cid)
#                 )
#
#         db.close()
#     except Exception as e:
#         print(f"Error starting countdown for streamer {streamer_id}: {e}")
#         import traceback
#         traceback.print_exc()


# async def stop_all_chests_countdown(streamer_id: int):
#     """停止主播所有宝箱的倒计时（已废弃）"""
#     try:
#         db = SessionLocal()
#         chests = GameService.get_active_chests(db, streamer_id)
#
#         for chest in chests:
#             await countdown_manager.stop_chest_countdown(chest.id)
#
#         db.close()
#     except Exception as e:
#         print(f"Error stopping countdown for streamer {streamer_id}: {e}")


# async def send_countdown_update(chest_id: int, time_remaining: int, streamer_id: int):
#     """发送倒计时更新消息（已废弃）"""
#     print(f"Sending countdown update: chest {chest_id}, time: {time_remaining}")
#     countdown_message = {
#         "type": "countdown_update",
#         "chest_id": chest_id,
#         "time_remaining": time_remaining
#     }
#
#     # 向主播和所有用户广播倒计时更新
#     await manager.send_to_streamer(streamer_id, countdown_message)
#     await manager.broadcast_to_all(countdown_message)


# async def handle_countdown_expire(chest_id: int):
#     """处理倒计时结束（已废弃）"""
#     print(f"Chest {chest_id} countdown expired, locking...")
#
#     try:
#         # 更新数据库中的宝箱状态
#         db = SessionLocal()
#         updated_chest = GameService.lock_expired_chest(db, chest_id)
#
#         if updated_chest:
#             # 通知所有客户端宝箱状态变更
#             status_message = {
#                 "type": "chest_status",
#                 "chest_id": chest_id,
#                 "status": updated_chest.status
#             }
#             await manager.broadcast_to_all(status_message)
#
#         db.close()
#     except Exception as e:
#         print(f"Error locking expired chest {chest_id}: {e}")
#         try:
#             db.close()
#         except:
#             pass


async def notify_user_balance_update(user_id: int, new_balance: int):
    """通知用户余额更新"""
    message = {
        "type": "balance_update",
        "balance": new_balance
    }
    await manager.send_to_user(user_id, message)