import asyncio
import json
import structlog
import time
from collections import defaultdict
from typing import Optional

from fastapi import WebSocket

logger = structlog.get_logger(__name__)


class ConnectionManager:
    def __init__(self):
        # user_id -> list of WebSocket connections
        self._connections: dict[int, list[WebSocket]] = defaultdict(list)
        # user_id -> set of friend user_ids (cached for broadcast)
        self._friend_cache: dict[int, set[int]] = {}
        self._cleanup_lock = asyncio.Lock()

    async def connect(self, user_id: int, ws: WebSocket):
        await ws.accept()
        self._connections[user_id].append(ws)
        logger.info("WS connected", user_id=user_id, total=len(self._connections))

    def disconnect(self, user_id: int, ws: WebSocket):
        conns = self._connections.get(user_id, [])
        if ws in conns:
            conns.remove(ws)
        if not self._connections.get(user_id):
            self._connections.pop(user_id, None)
        logger.info("WS disconnected", user_id=user_id, total=len(self._connections))

    async def send_json(self, user_id: int, data: dict):
        conns = self._connections.get(user_id, [])
        if not conns:
            return
        msg = json.dumps(data)
        dead = []
        for ws in conns:
            try:
                await ws.send_text(msg)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(user_id, ws)

    async def broadcast_friends_update(self, user_id: int, friends_data: list[dict]):
        """Broadcast friend list update to a specific user"""
        await self.send_json(user_id, {
            "type": "friends_update",
            "friends": friends_data
        })

    async def broadcast_status_change(self, user_id: int, friends: list[int]):
        """Notify all friends that this user's status changed"""
        data = {"type": "friend_status", "user_id": user_id}
        for fid in friends:
            await self.send_json(fid, data)

    async def broadcast_admin_clients(self, clients_data: list[dict], admin_ids: list[int]):
        """Broadcast client list to all connected admins"""
        data = {"type": "clients_update", "clients": clients_data}
        for aid in admin_ids:
            await self.send_json(aid, data)


manager = ConnectionManager()