from kaya.core import HttpContext, KayaApp, WebSocket
from kaya.session import InMemorySessionStore, SessionMixin

app = KayaApp(mixins=[SessionMixin(InMemorySessionStore())])


@app.GET('/')
async def home(ctx: HttpContext) -> None:
    visits = ctx.session.get('visits', 0) + 1
    ctx.session['visits'] = visits
    await ctx.send_str(200, f'visits: {visits}')


@app.GET('/read')
async def read(ctx: HttpContext) -> None:
    visits = ctx.session.get('visits', 0)
    await ctx.send_str(200, f'visits: {visits}')


@app.GET('/clear')
async def clear(ctx: HttpContext) -> None:
    ctx.session.invalidate()
    await ctx.send_str(200, 'session cleared')

@app.websocket('/echo')
async def echo(ws: WebSocket) -> None:
    await ws.accept()
    async for msg in ws:
        if msg.kind == 'text':
            await ws.send_text(f"echo: {msg.data}")
        elif msg.kind == 'binary':
            data = msg.data
            assert isinstance(data, bytes)
            await ws.send_bytes(data)

@app.websocket('/ws/visits')
async def ws_visits(ws: WebSocket) -> None:
    # WebSocket handlers can read the existing session. Most ASGI servers
    # (including Granian and Daphne) do not forward the `headers` field of the
    # `websocket.accept` message into the HTTP 101 response, so a new session
    # cookie cannot be set during the handshake. Use the HTTP `/` endpoint to
    # set or refresh the session cookie before connecting here.
    await ws.accept()
    visits = ws.session.get('visits', 0)
    await ws.send_text(f'visits: {visits}')