import asyncio
from typing import List
from uuid import UUID

from fastapi import WebSocket

PROJ_IOT_ONLINE_STATE = 0
PROJ_SENTINEL_VEHICLE_STATES = 1
PROJ_SENTINEL_MONITOR_STATUS = 2


class ConnectionManager:
    def __init__(self):
        self.active_connections: List[dict] = []  # 保存 websocket 和用户信息
        self.lock = asyncio.Lock()

    # proj_id:0：在线状态 1：畜牧车辆进入 2：畜牧监控大屏
    async def connect(
        self, websocket: WebSocket, user_id: UUID, dept_id: str, proj_id: int
    ):
        await websocket.accept()
        async with self.lock:
            self.active_connections.append(
                {
                    "ws": websocket,
                    "user_id": user_id,
                    "dept_id": dept_id,
                    "proj_id": proj_id,
                }
            )
            print(
                "Device" + str(user_id) + " in project " + str(proj_id) + " connected"
            )

    async def disconnect(self, websocket: WebSocket):
        async with self.lock:
            self.active_connections = [
                conn for conn in self.active_connections if conn["ws"] != websocket
            ]

    async def noticeOnlineStatus(self, message: dict):
        async with self.lock:
            for conn in self.active_connections:
                if conn["proj_id"] == PROJ_IOT_ONLINE_STATE:
                    await conn["ws"].send_json(message)

    async def noticeSentinel(
        self, message: dict, target_departments: List[UUID] = None
    ):
        """
        target_departments: 指定哪些部门能收到消息
        """
        async with self.lock:
            for conn in self.active_connections:
                if target_departments:
                    if (
                        conn["proj_id"] == PROJ_SENTINEL_VEHICLE_STATES
                        and conn["dept_id"] in target_departments
                    ):
                        await conn["ws"].send_json(message)

    async def noticeSentinelMonitorStatus(
        self, message: dict, target_departments: List[UUID] = None
    ):
        """
        target_departments: 指定哪些部门能收到消息
        """
        async with self.lock:
            for conn in self.active_connections:
                if target_departments:
                    if (
                        conn["proj_id"] == PROJ_SENTINEL_MONITOR_STATUS
                        and conn["dept_id"] in target_departments
                    ):
                        await conn["ws"].send_json(message)