import unittest
import httpx
from pwo import async_test
from httpx_ws import aconnect_ws, WebSocketDisconnect
from httpx_ws.transport import ASGIWebSocketTransport
from kaya.core import KayaApp, WebSocket


class WebSocketTest(unittest.TestCase):
    app: KayaApp

    def setUp(self):
        self.app = KayaApp()

        @self.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)

        @self.app.websocket('/room/${room_id}')
        async def room(ws: WebSocket, room_id: str) -> None:
            await ws.accept()
            async for msg in ws:
                if msg.kind == 'text':
                    await ws.send_text(f"[{room_id}] {msg.data}")

    @async_test
    async def test_echo_text(self):
        transport = ASGIWebSocketTransport(app=self.app)

        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
            async with aconnect_ws("/echo", client) as ws:
                await ws.send_text("hello")
                message = await ws.receive_text()
                self.assertEqual(message, "echo: hello")

    @async_test
    async def test_echo_binary(self):
        transport = ASGIWebSocketTransport(app=self.app)

        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
            async with aconnect_ws("/echo", client) as ws:
                await ws.send_bytes(b"hello")
                message = await ws.receive_bytes()
                self.assertEqual(message, b"hello")

    @async_test
    async def test_path_parameter(self):
        transport = ASGIWebSocketTransport(app=self.app)

        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
            async with aconnect_ws("/room/general", client) as ws:
                await ws.send_text("hi")
                message = await ws.receive_text()
                self.assertEqual(message, "[general] hi")

    @async_test
    async def test_no_handler(self):
        transport = ASGIWebSocketTransport(app=self.app)

        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
            with self.assertRaises(WebSocketDisconnect) as cm:
                async with aconnect_ws("/unknown", client) as ws:
                    pass

            self.assertEqual(cm.exception.code, 1000)

    @async_test
    async def test_close_from_client(self):
        transport = ASGIWebSocketTransport(app=self.app)

        async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as client:
            async with aconnect_ws("/echo", client) as ws:
                pass