from contextlib import asynccontextmanager
from typing import AsyncIterator

from fastapi import FastAPI
from fastapi_proxy_lib.core.websocket import (
    CallbackPipeContextType,
    ReverseWebSocketProxy,
)
from httpx import AsyncClient
from starlette.websockets import WebSocket

proxy = ReverseWebSocketProxy(AsyncClient(), base_url="ws://echo.websocket.events/")


async def client_to_server_callback(pipe_context: CallbackPipeContextType[str]) -> None:
    with pipe_context as (sender, receiver):
        async for message in receiver:
            print(f"Received from client: {message}")
            await sender.send(f"CTS:{message}")
    print("client_to_server_callback end")


async def server_to_client_callback(
    pipe_context: CallbackPipeContextType[str],
) -> None:
    with pipe_context as (sender, receiver):
        async for message in receiver:
            print(f"Received from server: {message}")
            await sender.send("STC:{message}")
    print("server_to_client_callback end")


@asynccontextmanager
async def close_proxy_event(_: FastAPI) -> AsyncIterator[None]:
    """Close proxy."""
    yield
    await proxy.aclose()


app = FastAPI(lifespan=close_proxy_event)


@app.websocket("/{path:path}")
async def _(websocket: WebSocket, path: str = ""):
    return await proxy.proxy(
        websocket=websocket,
        path=path,
        client_to_server_callback=client_to_server_callback,
        server_to_client_callback=server_to_client_callback,
    )


# Then run shell: `uvicorn <your_py>:app --host 127.0.0.1 --port 8000`
# visit the app: `ws://127.0.0.1:8000/`
# you can establish websocket connection with `ws://echo.websocket.events`

from httpx_ws import aconnect_ws

async with aconnect_ws('ws://127.0.0.1:8000/') as ws:
    message = await ws.receive_text()
    print(f"Received: {message}")

    await ws.send_text('Hello, World!')

    message = await ws.receive_text()
    print(f"Received: {message}")