from collections import defaultdict

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from sqlalchemy.orm import Session

from app.database import SessionLocal
from app.mongo import get_mongo_db
from app.services.document_service import get_document
from app.services.user_service import get_user_by_email
from app.utils.security import verify_token

router = APIRouter()

rooms: dict[str, set[WebSocket]] = defaultdict(set)
presence: dict[str, dict[int, str]] = defaultdict(dict)


async def _user_from_token(token: str):
    subject = verify_token(token, expected_type="access")
    if not subject:
        return None
    db: Session = SessionLocal()
    try:
        return get_user_by_email(db, subject)
    finally:
        db.close()


@router.websocket("/ws/collab/{doc_id}")
async def collab_socket(websocket: WebSocket, doc_id: str) -> None:
    token = websocket.query_params.get("token", "")
    user = await _user_from_token(token)
    if user is None:
        await websocket.close(code=4401)
        return

    try:
        mongo = get_mongo_db()
    except RuntimeError:
        await websocket.close(code=1013)
        return

    doc = await get_document(mongo, doc_id, user.id)
    if doc is None:
        await websocket.close(code=4403)
        return

    await websocket.accept()
    rooms[doc_id].add(websocket)
    presence[doc_id][user.id] = user.full_name or user.username

    async def broadcast(payload: dict, skip: WebSocket | None = None) -> None:
        dead: list[WebSocket] = []
        for client in rooms[doc_id]:
            if client is skip:
                continue
            try:
                await client.send_json(payload)
            except Exception:
                dead.append(client)
        for client in dead:
            rooms[doc_id].discard(client)

    await broadcast(
        {
            "type": "presence",
            "peers": [{"id": uid, "name": name} for uid, name in presence[doc_id].items()],
        }
    )

    try:
        while True:
            message = await websocket.receive_json()
            if message.get("type") == "content":
                await broadcast(
                    {
                        "type": "content",
                        "content": message.get("content"),
                        "from": user.id,
                    },
                    skip=websocket,
                )
    except WebSocketDisconnect:
        pass
    finally:
        rooms[doc_id].discard(websocket)
        presence[doc_id].pop(user.id, None)
        if rooms[doc_id]:
            await broadcast(
                {
                    "type": "presence",
                    "peers": [{"id": uid, "name": name} for uid, name in presence[doc_id].items()],
                }
            )
        else:
            rooms.pop(doc_id, None)
            presence.pop(doc_id, None)