import unittest
from types import SimpleNamespace
from typing import Any

from granian.rsgi import ProtocolClosed
from pwo import async_test

from kaya.core import WebSocketDisconnect
from kaya.rsgi import RsgiWebSocket


class FakeScope:
    scheme = 'ws'
    path = '/ws'
    query_string = ''
    headers: dict[str, list[str]] = {}
    client = '127.0.0.1:12345'
    server = '127.0.0.1:80'


class FakeTransport:
    """A transport whose peer vanished without a close handshake."""

    def __init__(self) -> None:
        self.receive_calls = 0
        self.sent: list[Any] = []

    async def receive(self) -> Any:
        self.receive_calls += 1
        raise ProtocolClosed()

    async def send_str(self, data: str) -> None:
        self.sent.append(data)
        raise ProtocolClosed()

    async def send_bytes(self, data: bytes) -> None:
        self.sent.append(data)
        raise ProtocolClosed()


class CleanCloseTransport(FakeTransport):
    """A transport that reports a close message (kind 0)."""

    async def receive(self) -> Any:
        self.receive_calls += 1
        return SimpleNamespace(kind=0, data=None)


class FakeProtocol:
    def __init__(self, transport: FakeTransport) -> None:
        self._transport = transport
        self.closed: list[int] = []

    async def accept(self) -> FakeTransport:
        return self._transport

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


class ClosedOnAcceptProtocol(FakeProtocol):
    """A protocol whose peer vanished during the handshake."""

    async def accept(self) -> FakeTransport:
        raise ProtocolClosed()


class RsgiWebSocketTest(unittest.TestCase):

    def test_misconfigured_granian(self):
        class FakeProtocolWithoutWebsockets:
            pass

        with self.assertRaises(RuntimeError) as ctx:
            RsgiWebSocket(FakeScope(), FakeProtocolWithoutWebsockets())  # type: ignore[arg-type]

        self.assertIn('Granian was not configured for websockets', str(ctx.exception))

    async def _socket(self, transport: FakeTransport) -> tuple[RsgiWebSocket, FakeProtocol]:
        protocol = FakeProtocol(transport)
        ws = RsgiWebSocket(FakeScope(), protocol)  # type: ignore[arg-type]
        await ws.accept()
        return ws, protocol

    @async_test
    async def test_receive_returns_close_when_transport_is_gone(self):
        ws, _ = await self._socket(FakeTransport())
        message = await ws.receive()
        self.assertEqual('close', message.kind)

    @async_test
    async def test_iteration_ends_when_transport_is_gone(self):
        ws, _ = await self._socket(FakeTransport())
        messages = [message async for message in ws]
        self.assertEqual([], messages)

    @async_test
    async def test_receive_remembers_the_close(self):
        transport = FakeTransport()
        ws, _ = await self._socket(transport)
        self.assertEqual('close', (await ws.receive()).kind)
        self.assertEqual('close', (await ws.receive()).kind)
        self.assertEqual(1, transport.receive_calls)

    @async_test
    async def test_send_text_raises_disconnect_when_transport_is_gone(self):
        ws, _ = await self._socket(FakeTransport())
        with self.assertRaises(WebSocketDisconnect) as ctx:
            await ws.send_text('hello')
        self.assertEqual(1006, ctx.exception.code)

    @async_test
    async def test_send_bytes_raises_disconnect_when_transport_is_gone(self):
        ws, _ = await self._socket(FakeTransport())
        with self.assertRaises(WebSocketDisconnect) as ctx:
            await ws.send_bytes(b'hello')
        self.assertEqual(1006, ctx.exception.code)

    @async_test
    async def test_sends_after_disconnect_do_not_touch_the_transport(self):
        transport = FakeTransport()
        ws, _ = await self._socket(transport)
        with self.assertRaises(WebSocketDisconnect):
            await ws.send_text('first')
        with self.assertRaises(WebSocketDisconnect):
            await ws.send_text('second')
        self.assertEqual(['first'], transport.sent)

    @async_test
    async def test_send_after_a_clean_close_raises_disconnect(self):
        transport = CleanCloseTransport()
        ws, _ = await self._socket(transport)
        self.assertEqual('close', (await ws.receive()).kind)
        with self.assertRaises(WebSocketDisconnect):
            await ws.send_text('hello')
        self.assertEqual([], transport.sent)

    @async_test
    async def test_accept_raises_disconnect_when_transport_is_gone(self):
        ws = RsgiWebSocket(FakeScope(), ClosedOnAcceptProtocol(FakeTransport()))  # type: ignore[arg-type]
        with self.assertRaises(WebSocketDisconnect) as ctx:
            await ws.accept()
        self.assertEqual(1006, ctx.exception.code)

    @async_test
    async def test_close_is_noop_once_the_transport_is_gone(self):
        ws, protocol = await self._socket(FakeTransport())
        await ws.receive()
        await ws.close(1000)
        self.assertEqual([], protocol.closed)