from __future__ import annotations

import asyncio
from dataclasses import dataclass
from starlette.websockets import WebSocket


@dataclass
class ConnectedUser:
    """Minimal user info stored per WebSocket connection."""
    user_id: str
    role: str  # admin / user / read-only / pending


class ConnectionManager:
    """Track active WebSocket connections and broadcast events."""

    def __init__(self, max_connections: int = 50) -> None:
        self._connections: dict[WebSocket, ConnectedUser] = {}
        self._thread_owners: dict[str, str] = {}  # thread_id → user_id
        self._max = max_connections
        self._lock = asyncio.Lock()

    async def accept(self, ws: WebSocket, user_id: str, role: str = "user") -> bool:
        """Add connection if under limit. Returns False if over limit."""
        async with self._lock:
            if len(self._connections) >= self._max:
                return False
            self._connections[ws] = ConnectedUser(user_id=user_id, role=role)
            return True

    def disconnect(self, ws: WebSocket) -> None:
        """Remove connection from registry."""
        self._connections.pop(ws, None)

    def get_user_id(self, ws: WebSocket) -> str | None:
        """Return the user_id associated with a WebSocket, or None."""
        cu = self._connections.get(ws)
        return cu.user_id if cu else None

    def get_role(self, ws: WebSocket) -> str | None:
        """Return the role associated with a WebSocket, or None."""
        cu = self._connections.get(ws)
        return cu.role if cu else None

    async def broadcast(self, data: dict) -> None:
        """Send data to all connected clients. Removes dead connections."""
        dead = []
        for ws in list(self._connections):
            try:
                await ws.send_json(data)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self._connections.pop(ws, None)

    async def broadcast_to_user(self, user_id: str, data: dict) -> None:
        """Send data only to connections belonging to a specific user."""
        dead = []
        for ws, cu in list(self._connections.items()):
            if cu.user_id != user_id:
                continue
            try:
                await ws.send_json(data)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self._connections.pop(ws, None)

    async def send_to(self, ws: WebSocket, data: dict) -> None:
        """Send data to one specific connection."""
        try:
            await ws.send_json(data)
        except Exception:
            self._connections.pop(ws, None)

    # ── Thread ownership tracking (for scoped broadcasts) ───

    def register_thread(self, thread_id: str, user_id: str) -> None:
        """Track which user owns an active thread."""
        self._thread_owners[thread_id] = user_id

    def unregister_thread(self, thread_id: str) -> None:
        """Remove thread ownership tracking."""
        self._thread_owners.pop(thread_id, None)

    def get_thread_owner(self, thread_id: str) -> str | None:
        """Return the user_id that owns the given thread, or None."""
        return self._thread_owners.get(thread_id)

    @property
    def connection_count(self) -> int:
        return len(self._connections)
