import unittest
import httpx
from pwo import async_test
from httpx_ws import aconnect_ws, WebSocketDisconnect as HttpxWebSocketDisconnect
from httpx_ws.transport import ASGIWebSocketTransport
from kaya.core import KayaApp, WebSocket, WebSocketDisconnect, WebSocketMessage
from kaya.core._asgi import AsgiWebSocket


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(HttpxWebSocketDisconnect) 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

    @async_test
    async def test_accept_with_headers(self):
        sent_messages = []

        async def send(message):
            sent_messages.append(message)

        async def receive():
            return {'type': 'websocket.connect'}

        scope = {
            'type': 'websocket',
            'path': '/echo',
            'query_string': b'',
            'scheme': 'ws',
            'client': ('127.0.0.1', 12345),
            'server': ('127.0.0.1', 80),
            'headers': [],
        }
        ws = AsgiWebSocket(scope, receive, send)
        await ws.accept(headers={'Set-Cookie': 'sid=abc; Path=/', 'X-Custom': ('a', 'b')})
        self.assertEqual(1, len(sent_messages))
        message = sent_messages[0]
        self.assertEqual('websocket.accept', message['type'])
        self.assertIn((b'Set-Cookie', b'sid=abc; Path=/'), message['headers'])
        self.assertIn((b'X-Custom', b'a'), message['headers'])
        self.assertIn((b'X-Custom', b'b'), message['headers'])

    @async_test
    async def test_accept_without_headers(self):
        sent_messages = []

        async def send(message):
            sent_messages.append(message)

        async def receive():
            return {'type': 'websocket.connect'}

        scope = {
            'type': 'websocket',
            'path': '/echo',
            'query_string': b'',
            'scheme': 'ws',
            'client': ('127.0.0.1', 12345),
            'server': ('127.0.0.1', 80),
            'headers': [],
        }
        ws = AsgiWebSocket(scope, receive, send)
        await ws.accept()
        self.assertEqual(1, len(sent_messages))
        self.assertEqual({'type': 'websocket.accept'}, sent_messages[0])

    @async_test
    async def test_websocket_scope_without_scheme(self):
        # Daphne omits the optional `scheme` key from websocket scopes.
        async def send(message):
            pass

        async def receive():
            return {'type': 'websocket.connect'}

        scope = {
            'type': 'websocket',
            'path': '/echo',
            'query_string': b'',
            'client': ('127.0.0.1', 12345),
            'server': ('127.0.0.1', 80),
            'headers': [],
        }
        ws = AsgiWebSocket(scope, receive, send)
        self.assertEqual('ws', ws.scheme)


class FakeWebSocket(WebSocket):
    """Minimal in-memory websocket for exercising ``KayaApp`` directly."""

    def __init__(self, path: str) -> None:
        self.path = path
        self.query_string = ''
        self.scheme = 'ws'
        self.headers: dict[str, list[str]] = {}
        self.client = ('127.0.0.1', 12345)
        self.server = ('127.0.0.1', 80)
        self.closed: list[int] = []

    async def accept(self, headers=None) -> None:
        pass

    async def receive(self) -> WebSocketMessage:
        return WebSocketMessage(kind='close')

    async def send_text(self, data: str) -> None:
        pass

    async def send_bytes(self, data: bytes) -> None:
        pass

    async def close(self, code: int = 1000) -> None:
        self.closed.append(code)

    async def __anext__(self) -> WebSocketMessage:
        raise StopAsyncIteration


class WebSocketDisconnectTest(unittest.TestCase):

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

        @self.app.websocket('/disconnect')
        async def disconnect(ws: WebSocket) -> None:
            await ws.accept()
            await ws.send_text('before')
            raise WebSocketDisconnect()

        @self.app.websocket('/failure')
        async def failure(ws: WebSocket) -> None:
            await ws.accept()
            raise RuntimeError('boom')

    @async_test
    async def test_disconnect_terminates_quietly(self):
        after_hook_calls = []

        async def after(ws: WebSocket) -> None:
            after_hook_calls.append(ws)

        self.app.add_after_websocket_hook(after)
        ws = FakeWebSocket('/disconnect')
        await self.app.handle_websocket(ws)
        self.assertEqual([ws], after_hook_calls)
        self.assertIsNone(ws.exception)

    @async_test
    async def test_other_errors_are_recorded_and_reraised(self):
        ws = FakeWebSocket('/failure')
        with self.assertRaises(RuntimeError):
            await self.app.handle_websocket(ws)
        self.assertIsInstance(ws.exception, RuntimeError)

    @async_test
    async def test_asgi_send_after_disconnect_raises_websocket_disconnect(self):
        async def send(message):
            raise OSError('client gone')

        async def receive():
            return {'type': 'websocket.disconnect', 'code': 1006}

        ws = AsgiWebSocket(self._scope(), receive, send)
        with self.assertRaises(WebSocketDisconnect) as ctx:
            await ws.send_text('hello')
        self.assertEqual(1006, ctx.exception.code)
        with self.assertRaises(WebSocketDisconnect):
            await ws.send_bytes(b'hello')
        with self.assertRaises(WebSocketDisconnect):
            await ws.close()

    @async_test
    async def test_asgi_accept_after_disconnect_raises_websocket_disconnect(self):
        async def send(message):
            raise OSError('client gone')

        async def receive():
            return {'type': 'websocket.connect'}

        ws = AsgiWebSocket(self._scope(), receive, send)
        with self.assertRaises(WebSocketDisconnect) as ctx:
            await ws.accept()
        self.assertEqual(1006, ctx.exception.code)

    @staticmethod
    def _scope() -> dict:
        return {
            'type': 'websocket',
            'path': '/echo',
            'query_string': b'',
            'scheme': 'ws',
            'client': ('127.0.0.1', 12345),
            'server': ('127.0.0.1', 80),
            'headers': [],
        }