from fastapi import APIRouter, Query
from starlette.websockets import WebSocket, WebSocketDisconnect

from config.security import get_user_id_from_token_from_ws
from db.postgres import get_dept_id_by_user_id
from db.postgres.ws_manager import ConnectionManager

ws_manager = ConnectionManager()


iot_ws_router = APIRouter()


# EMQX 通知在线与否
@iot_ws_router.websocket("/device-status")
async def websocket_device_status(
    websocket: WebSocket,
    token: str = Query(...),
):
    user_id = get_user_id_from_token_from_ws(token)
    dept_id = get_dept_id_by_user_id(user_id)  # 查数据库或缓存
    await ws_manager.connect(websocket, user_id, dept_id, 0)
    print("[WS] client connected")

    try:
        while True:
            # 这里不需要接收任何消息
            await websocket.receive_text()
    except WebSocketDisconnect:
        await ws_manager.disconnect(websocket)
        print("[WS] client disconnected")


# Vue 牧安云哨
@iot_ws_router.websocket("/sentinel_record")
async def websocket_sentinel_record(
    websocket: WebSocket,
    token: str = Query(...),
):
    user_id = get_user_id_from_token_from_ws(token)
    dept_id = get_dept_id_by_user_id(user_id)  # 查数据库或缓存
    print("user_id:", user_id)
    print("dept_id:", dept_id)
    print("已接入")
    await ws_manager.connect(websocket, user_id, dept_id, 1)

    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        await ws_manager.disconnect(websocket)