{# templates/wspipe.py.jinja2 #}
import asyncio
import urllib.parse

from rich import print as rprint
from superagentx.pipeimpl.wspipe import WSPipe  # https://websockets.readthedocs.io/en/stable/
from websockets import CloseCode, ServerConnection

from {{ package_name }}.config import AUTH_TOKEN
from {{ package_name }}.pipe import get_{{ pipe_name }}_pipe


async def auth_handler(ws: ServerConnection) -> bool:
    """Authenticate user from token in query parameter."""
    query = urllib.parse.urlparse(ws.request.path).query
    params = urllib.parse.parse_qs(query)
    values = params.get('token', [])
    if values:
        token = values[0]
        if token is None or token != AUTH_TOKEN:
            await ws.close(
                code=CloseCode.INTERNAL_ERROR,
                reason='Authentication failed!'
            )
            return False
    return True


async def main():
    """
    Launches the {{ app_name }} pipeline websocket server for processing requests and handling data.
    """
    pipe = await get_{{ pipe_name }}_pipe()
    ws_pipe = WSPipe(
        search_name='SuperAgentX {{ app_name }} Websocket Server',
        agentx_pipe=pipe,
        auth_handler=auth_handler,
    )
    await ws_pipe.start()


if __name__ == '__main__':
    try:
        asyncio.run(main())
    except (KeyboardInterrupt, asyncio.CancelledError):
        rprint("\nUser canceled the [bold yellow][i]pipe[/i]!"){{'\n'}}
