mtsd/backend/app/ws/price_updates.py

"""WebSocket：区域价格清单更新推送（客户端实时刷新）"""

from __future__ import annotations

import asyncio
import logging
from dataclasses import dataclass
from typing import Dict, Set, Optional, Any

import anyio
from fastapi import WebSocket, WebSocketDisconnect

from app.core.database import SessionLocal
from app.core.auth import get_current_user_by_token

logger = logging.getLogger(__name__)

@dataclass(frozen=True)
class PriceListUpdatedEvent:
    regionId: str
    priceListId: str
    priceVersion: int
    publishedTime: Optional[str] = None

    def to_message(self) -> dict:
        return {
            "event": "price_list_updated",
            "data": {
                "regionId": self.regionId,
                "priceListId": self.priceListId,
                "priceVersion": self.priceVersion,
                "publishedTime": self.publishedTime,
            },
        }


class RegionWSManager:
    """按 region_id 管理 websocket 连接并广播。"""

    def __init__(self) -> None:
        self._lock = asyncio.Lock()
        self._rooms: Dict[str, Set[WebSocket]] = {}

    async def connect(self, websocket: WebSocket, *, token: str) -> dict:
        """鉴权并接入房间。返回 current_user dict。"""
        # 验证 token（包含 Redis 会话校验 + DB 用户校验）
        db = SessionLocal()
        try:
            current_user = await get_current_user_by_token(token, db)
        finally:
            db.close()

        region_id = (current_user.get("region_id") or "").strip()
        if not region_id:
            logger.warning("[ws-mini] reject: empty region_id user=%s", current_user.get("id"))
            await websocket.close(code=1008)
            raise RuntimeError("user has no region_id")

        await websocket.accept()

        async with self._lock:
            self._rooms.setdefault(region_id, set()).add(websocket)
            room_size = len(self._rooms.get(region_id, set()))
        logger.warning(
            "[ws-mini] connected user=%s role=%s region=%s room_size=%s",
            current_user.get("id"),
            current_user.get("role"),
            region_id,
            room_size,
        )

        return current_user

    async def disconnect(self, websocket: WebSocket, *, region_id: str) -> None:
        async with self._lock:
            s = self._rooms.get(region_id)
            if not s:
                return
            if websocket in s:
                s.remove(websocket)
            room_size = len(s)
            if not s:
                self._rooms.pop(region_id, None)
        logger.warning("[ws-mini] disconnected region=%s room_size=%s", region_id, room_size if "room_size" in locals() else 0)

    async def broadcast_region(self, region_id: str, message: dict) -> int:
        """向 region 房间广播消息。返回尝试发送的连接数。"""
        async with self._lock:
            targets = list(self._rooms.get(region_id, set()))

        if not targets:
            logger.warning("[ws-mini] broadcast skip: no targets region=%s event=%s", region_id, message.get("event"))
            return 0

        logger.warning(
            "[ws-mini] broadcasting region=%s event=%s target_count=%s",
            region_id,
            message.get("event"),
            len(targets),
        )
        dead: list[WebSocket] = []
        for ws in targets:
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)

        if dead:
            async with self._lock:
                room = self._rooms.get(region_id)
                if room:
                    for ws in dead:
                        room.discard(ws)
                    if not room:
                        self._rooms.pop(region_id, None)

        return len(targets)


manager = RegionWSManager()


async def ws_mini_price_updates(websocket: WebSocket) -> None:
    """小程序端 WebSocket 入口：/ws/mini"""
    token = (websocket.query_params.get("token") or "").strip()
    if not token:
        # 兼容 Authorization: Bearer <token>
        auth = (websocket.headers.get("authorization") or "").strip()
        if auth.lower().startswith("bearer "):
            token = auth[7:].strip()

    if not token:
        logger.warning("[ws-mini] reject: missing token ip=%s", websocket.client.host if websocket.client else "")
        await websocket.close(code=1008)
        return

    try:
        current_user = await manager.connect(websocket, token=token)
    except Exception:
        logger.exception("[ws-mini] connect failed")
        # connect 内部会 close；这里兜底
        try:
            await websocket.close(code=1008)
        except Exception:
            pass
        return

    region_id = (current_user.get("region_id") or "").strip()
    try:
        # 当前实现：仅服务器推送；客户端无需主动发消息。
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        logger.warning("[ws-mini] client disconnected user=%s region=%s", current_user.get("id"), region_id)
        pass
    except Exception:
        logger.exception("[ws-mini] receive loop error user=%s region=%s", current_user.get("id"), region_id)
        pass
    finally:
        await manager.disconnect(websocket, region_id=region_id)


def notify_price_list_updated(event: PriceListUpdatedEvent) -> None:
    """线程/同步上下文安全：广播价格更新事件。"""

    async def _run():
        sent = await manager.broadcast_region(event.regionId, event.to_message())
        logger.warning(
            "[ws-mini] notify price update region=%s price_list_id=%s version=%s sent=%s",
            event.regionId,
            event.priceListId,
            event.priceVersion,
            sent,
        )

    try:
        loop = asyncio.get_running_loop()
        loop.create_task(_run())
    except RuntimeError:
        # sync endpoint 常在 worker thread：使用 anyio 跳回事件循环
        anyio.from_thread.run(_run)


def register_ws_routes(app: Any) -> None:
    """在 FastAPI app 上注册 websocket 路由。"""
    app.websocket("/ws/mini")(ws_mini_price_updates)