from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from n6k_server.session import (
    SESSION_COOKIE,
    create_session,
    destroy_session,
)

router = APIRouter()


@router.websocket("/ws")
async def ws_session(websocket: WebSocket):
    app = websocket.app
    registry = app.state.table_registry
    loaded_views = {n: s.sql for n, s in registry._views.items() if s.registered}
    session = create_session(app.state.cached_tables, loaded_views)
    app.state.sessions[session.id] = session

    cookie = f"{SESSION_COOKIE}={session.id}; HttpOnly; SameSite=Lax; Path=/"
    await websocket.accept(headers=[(b"set-cookie", cookie.encode())])

    await websocket.send_json({"session_id": session.id})

    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    finally:
        app.state.sessions.pop(session.id, None)
        destroy_session(session)
