"""WebSocket 端点：实时通知 + 文献推送"""

import logging

from fastapi import APIRouter, Query, WebSocket, WebSocketDisconnect

from app.core.security import decode_token
from app.core.tenant_context import tenant_ctx
from app.core.websocket import ws_manager

logger = logging.getLogger(__name__)

router = APIRouter()


@router.websocket("/feed")
async def websocket_feed(ws: WebSocket, token: str = Query(...)):
    """WebSocket 连接：认证后自动推送文献更新"""
    user_id = None
    try:
        payload = decode_token(token)
        if payload.get("type") != "access":
            await ws.close(code=4001)
            return
        user_id = payload.get("sub")
        if not user_id:
            await ws.close(code=4001)
            return
        # 设置租户上下文
        tid = payload.get("tid")
        if tid:
            tenant_ctx.set(tid)
    except Exception:
        await ws.close(code=4001)
        return

    await ws_manager.connect(user_id, ws)

    try:
        # 发送连接成功通知
        await ws.send_json({
            "type": "connected",
            "user_id": user_id,
            "online_users": ws_manager.online_users,
            "message": "实时推送已连接",
        })

        # 保持连接 + 心跳
        while True:
            data = await ws.receive_text()
            # 客户端可以发送 ping
            if data == "ping":
                await ws.send_json({"type": "pong"})

    except WebSocketDisconnect:
        pass
    except Exception:
        logger.warning("WS unexpected error for user %s", user_id, exc_info=True)
    finally:
        if user_id:
            await ws_manager.disconnect(user_id, ws)


@router.get("/status", summary="连接状态")
async def ws_status():
    """WebSocket 服务状态"""
    return {
        "online_users": ws_manager.online_users,
        "online_connections": ws_manager.online_count,
    }
