import asyncio
from uuid import UUID

from fastapi import WebSocket


class ConsolePresenceManager:
    def __init__(self) -> None:
        self._connections: dict[UUID, WebSocket] = {}
        self._relay_queues: dict[UUID, asyncio.Queue[str]] = {}
        self._lock = asyncio.Lock()

    async def register(self, user_id: UUID, websocket: WebSocket) -> None:
        async with self._lock:
            self._connections[user_id] = websocket

    async def unregister(self, user_id: UUID, websocket: WebSocket) -> None:
        async with self._lock:
            if self._connections.get(user_id) is websocket:
                self._connections.pop(user_id, None)
                self._relay_queues.pop(user_id, None)

    async def is_present(self, user_id: UUID) -> bool:
        async with self._lock:
            return user_id in self._connections

    async def get_connection(self, user_id: UUID) -> WebSocket | None:
        async with self._lock:
            return self._connections.get(user_id)

    async def start_relay(self, user_id: UUID) -> asyncio.Queue[str]:
        async with self._lock:
            queue: asyncio.Queue[str] = asyncio.Queue()
            self._relay_queues[user_id] = queue
            return queue

    async def stop_relay(self, user_id: UUID) -> None:
        async with self._lock:
            self._relay_queues.pop(user_id, None)

    async def forward_from_console(self, user_id: UUID, text: str) -> bool:
        async with self._lock:
            queue = self._relay_queues.get(user_id)
        if not queue:
            return False
        await queue.put(text)
        return True


console_presence = ConsolePresenceManager()