from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from sqlmodel import Session
from websockets.exceptions import ConnectionClosed
from learn_sql_model.api.websocket_connection_manager import manager
from learn_sql_model.config import get_session
from learn_sql_model.models.hero import HeroUpdate, Heros
web_socket_router = APIRouter()
html = """
Chat
WebSocket Chat
"""
@web_socket_router.get("/watch")
async def get():
return HTMLResponse(html)
@web_socket_router.websocket("/ws/{id}")
async def websocket_endpoint_connect(websocket: WebSocket, id: int):
await manager.connect(websocket, id)
try:
while True:
data = await websocket.receive_text()
await websocket.send_text(f"[gold]You Said: {data}")
await manager.broadcast(f"[blue]USER: {data}", id)
except WebSocketDisconnect:
manager.disconnect(websocket, id)
await manager.broadcast(f"Client #{id} left the chat", id)
@web_socket_router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
await manager.connect(websocket)
try:
while True:
data = await websocket.receive_text()
await manager.broadcast(f"[blue]USER: {data}")
except WebSocketDisconnect:
manager.disconnect(websocket, id)
await manager.broadcast(f"Client #{id} left the chat", id)
@web_socket_router.websocket("/wsecho")
async def websocket_endpoint(
websocket: WebSocket,
session: Session = Depends(get_session),
):
await websocket.accept()
try:
while True:
heros = Heros.list(session=session)
await websocket.send_text(heros.json())
except WebSocketDisconnect:
print("disconnected")
except ConnectionClosed:
print("connection closed")
@web_socket_router.websocket("/ws-hero-update")
async def websocket_endpoint(
websocket: WebSocket,
session: Session = Depends(get_session),
):
await websocket.accept()
try:
while True:
data = await websocket.receive_text()
hero = HeroUpdate.parse_raw(data)
print(hero)
hero.update(session=session)
print("hero is updated")
except WebSocketDisconnect:
print("disconnected")
except ConnectionClosed:
print("connection closed")