backend/api/session_manager.py

from __future__ import annotations
from fastapi import WebSocket


class SessionManager:
    def __init__(self):
        self._user_to_ws: dict[str, WebSocket] = {}
        self._ws_to_user: dict[str, str] = {}

    def register(self, user_id: str, ws_id: str, ws: WebSocket) -> None:
        # Remove stale reverse-map entry for the previous ws_id of this user
        old_ws_id = next((k for k, v in self._ws_to_user.items() if v == user_id), None)
        if old_ws_id:
            self._ws_to_user.pop(old_ws_id)
        self._user_to_ws[user_id] = ws
        self._ws_to_user[ws_id] = user_id

    async def register_and_evict(self, user_id: str, ws_id: str, ws: WebSocket) -> bool:
        old_ws = self._user_to_ws.get(user_id)
        evicted = False
        if old_ws and old_ws is not ws:
            try:
                await old_ws.send_json({"type": "session_evicted", "reason": "logged in elsewhere"})
                await old_ws.close()
            except Exception:
                pass
            evicted = True
        self.register(user_id, ws_id, ws)
        return evicted

    def unregister(self, user_id: str) -> None:
        ws = self._user_to_ws.pop(user_id, None)
        if ws:
            to_remove = [k for k, v in self._ws_to_user.items() if v == user_id]
            for k in to_remove:
                self._ws_to_user.pop(k, None)

    def get_ws(self, user_id: str) -> WebSocket | None:
        return self._user_to_ws.get(user_id)

    def get_user_id(self, ws_id: str) -> str | None:
        return self._ws_to_user.get(ws_id)

    def online_user_ids(self) -> list[str]:
        return list(self._user_to_ws.keys())