# Requires: `starlette`, `uvicorn`, `jinja2`
# Run with `uvicorn example:app`
import anyio
from broadcaster import Broadcast
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi import APIRouter, Response
import os
import logging
import asyncio

log = logging.getLogger(__name__)
router = APIRouter()


brd = Broadcast("memory://")


@router.get("/ws2_example")
async def get() -> HTMLResponse:

    html = open(f"{os.getcwd()}/webui/routers/messenger.html").read()
    return HTMLResponse(html)

@router.websocket("/ws2")
async def chatroom_ws(websocket: WebSocket) -> None:
    await websocket.accept()


    async with anyio.create_task_group() as task_group:
        # run until first is complete
        async def run_chatroom_ws_receiver() -> None:
            await chatroom_ws_receiver(websocket=websocket)
            task_group.cancel_scope.cancel()

        task_group.start_soon(run_chatroom_ws_receiver)
        log.warning("Started chatroom_ws_sender")

        await chatroom_ws_sender(websocket)


async def chatroom_ws_receiver(websocket: WebSocket) -> None:
    async for message in websocket.iter_text():
        log.warning(f"Received message: {message}")
        await brd.publish(channel="chatroom", message=message)


async def chatroom_ws_sender(websocket: WebSocket) -> None:
    async with brd.subscribe(channel="chatroom") as subscriber:
        log.warning("====>Subscribed to chatroom channel")
        async for event in subscriber:
            log.warning(f"Sending message: {event.message}")
            await websocket.send_text(event.message)