import asyncio
import json

import pytest
import websockets

from unittest.mock import AsyncMock, MagicMock, patch

from mrp.auth import Keypair
from mrp.websocket import WSClient, ReconnectingWSClient


class TestWSUrl:
    def test_http_to_ws(self):
        ws = WSClient("http://localhost:8080", Keypair.generate())
        assert ws._ws_url() == "ws://localhost:8080/v1/ws"

    def test_https_to_wss(self):
        ws = WSClient("https://relay.example.com", Keypair.generate())
        assert ws._ws_url() == "wss://relay.example.com/v1/ws"

    def test_trailing_slash_stripped(self):
        ws = WSClient("http://localhost:8080/", Keypair.generate())
        assert ws._ws_url() == "ws://localhost:8080/v1/ws"


def _make_mock_conn(recv_queue: asyncio.Queue):
    """Create a mock WebSocket connection backed by an asyncio.Queue."""
    mock_conn = AsyncMock()
    mock_conn.send = AsyncMock()
    mock_conn.close = AsyncMock()

    async def mock_recv():
        return await recv_queue.get()

    mock_conn.recv = mock_recv

    async def mock_aiter(self):
        while True:
            try:
                data = await asyncio.wait_for(recv_queue.get(), timeout=0.5)
                yield data
            except asyncio.TimeoutError:
                break

    mock_conn.__aiter__ = mock_aiter
    return mock_conn


class TestWSConnect:
    @pytest.mark.asyncio
    async def test_auth_handshake(self):
        """Verify the auth frame is sent with required fields and auth_result is consumed."""
        recv_queue = asyncio.Queue()
        mock_conn = _make_mock_conn(recv_queue)

        # Queue the auth_result response
        await recv_queue.put(json.dumps({"type": "auth_result", "status": "ok"}))

        keypair = Keypair.generate()
        ws = WSClient("http://localhost:8080", keypair)

        mock_connect = AsyncMock(return_value=mock_conn)
        with patch("mrp.websocket.websockets.connect", mock_connect):
            await ws.connect()

        # Verify auth frame was sent
        assert mock_conn.send.call_count == 1
        auth_frame = json.loads(mock_conn.send.call_args[0][0])
        assert auth_frame["type"] == "auth"
        assert auth_frame["public_key"] == keypair.public_key_b64
        assert "timestamp" in auth_frame
        assert "signature" in auth_frame

    @pytest.mark.asyncio
    async def test_auth_failure(self):
        """Server returns auth error -- should raise ConnectionError."""
        recv_queue = asyncio.Queue()
        mock_conn = _make_mock_conn(recv_queue)

        await recv_queue.put(json.dumps({
            "type": "auth_result", "status": "error", "error": "bad"
        }))

        keypair = Keypair.generate()
        ws = WSClient("http://localhost:8080", keypair)

        mock_connect = AsyncMock(return_value=mock_conn)
        with patch("mrp.websocket.websockets.connect", mock_connect):
            with pytest.raises(ConnectionError, match="Auth failed"):
                await ws.connect()

    @pytest.mark.asyncio
    async def test_auth_timeout(self):
        """recv never returns -- should raise ConnectionError after timeout."""
        recv_queue = asyncio.Queue()
        mock_conn = _make_mock_conn(recv_queue)

        # Override recv to hang
        async def slow_recv():
            await asyncio.sleep(60)

        mock_conn.recv = slow_recv

        keypair = Keypair.generate()
        ws = WSClient("http://localhost:8080", keypair)

        mock_connect = AsyncMock(return_value=mock_conn)
        with patch("mrp.websocket.websockets.connect", mock_connect):
            # Patch the timeout to be very short so the test is fast
            with patch("mrp.websocket.asyncio.wait_for", side_effect=asyncio.TimeoutError):
                with pytest.raises(ConnectionError, match="Auth response timeout"):
                    await ws.connect()


class TestWSMessageLoop:
    @pytest.mark.asyncio
    async def test_ping_pong(self):
        """Put a ping frame in the queue and verify pong is sent back."""
        recv_queue = asyncio.Queue()
        mock_conn = _make_mock_conn(recv_queue)

        await recv_queue.put(json.dumps({"type": "ping"}))

        keypair = Keypair.generate()
        ws = WSClient("http://localhost:8080", keypair)
        ws._ws = mock_conn

        await ws._read_loop()

        # Find the pong send call
        pong_calls = [
            call for call in mock_conn.send.call_args_list
            if json.loads(call[0][0]).get("type") == "pong"
        ]
        assert len(pong_calls) == 1

    @pytest.mark.asyncio
    async def test_receive_message_and_ack(self):
        """Put a message frame in the queue, verify on_message called and ack sent."""
        recv_queue = asyncio.Queue()
        mock_conn = _make_mock_conn(recv_queue)

        msg_frame = {
            "type": "message",
            "message": {
                "message_id": "msg_1",
                "sender_key": "sender_abc",
                "recipient_key": "recipient_xyz",
                "content_type": "application/json",
                "body": {"text": "hello"},
                "attachments": [],
                "status": "delivered",
                "metadata": {},
                "created_at": "2026-03-05T12:00:00Z",
                "expires_at": "2026-03-12T12:00:00Z",
            },
        }
        await recv_queue.put(json.dumps(msg_frame))

        keypair = Keypair.generate()
        ws = WSClient("http://localhost:8080", keypair)
        ws._ws = mock_conn

        received_messages = []

        async def handler(msg):
            received_messages.append(msg)
            return None

        ws.on_message(handler)
        await ws._read_loop()

        # Verify handler was called
        assert len(received_messages) == 1
        assert received_messages[0].message_id == "msg_1"
        assert received_messages[0].body == {"text": "hello"}

        # Verify ack was sent
        ack_calls = [
            call for call in mock_conn.send.call_args_list
            if json.loads(call[0][0]).get("type") == "ack"
        ]
        assert len(ack_calls) == 1
        ack_frame = json.loads(ack_calls[0][0][0])
        assert ack_frame["message_id"] == "msg_1"

    @pytest.mark.asyncio
    async def test_send_result_resolves(self):
        """Create a pending send, put send_result in queue, verify future resolved."""
        recv_queue = asyncio.Queue()
        mock_conn = _make_mock_conn(recv_queue)

        request_id = "req_123"
        send_result_frame = {
            "type": "send_result",
            "request_id": request_id,
            "status": "ok",
            "message_id": "msg_99",
        }
        await recv_queue.put(json.dumps(send_result_frame))

        keypair = Keypair.generate()
        ws = WSClient("http://localhost:8080", keypair)
        ws._ws = mock_conn

        # Create a pending future
        loop = asyncio.get_event_loop()
        future = loop.create_future()
        ws._pending_sends[request_id] = future

        await ws._read_loop()

        # Future should be resolved
        assert future.done()
        result = future.result()
        assert result["status"] == "ok"
        assert result["message_id"] == "msg_99"


class TestWSSendMessage:
    @pytest.mark.asyncio
    async def test_send_success(self):
        """Mock ws connected, call send_message, resolve the pending future, verify frame."""
        recv_queue = asyncio.Queue()
        mock_conn = _make_mock_conn(recv_queue)
        mock_conn.send = AsyncMock()

        keypair = Keypair.generate()
        ws = WSClient("http://localhost:8080", keypair)
        ws._ws = mock_conn

        async def resolve_send():
            # Wait for the pending send to be registered
            while not ws._pending_sends:
                await asyncio.sleep(0.01)
            # Resolve it
            for req_id, future in ws._pending_sends.items():
                if not future.done():
                    future.set_result({
                        "status": "ok",
                        "message_id": "msg_42",
                        "request_id": req_id,
                    })

        # Run send_message and the resolver concurrently
        result, _ = await asyncio.gather(
            ws.send_message("recipient_key_abc", {"text": "hi"}),
            resolve_send(),
        )

        assert result.message_id == "msg_42"
        assert result.status == "sent"

        # Verify the sent frame structure
        sent_data = json.loads(mock_conn.send.call_args[0][0])
        assert sent_data["type"] == "send"
        assert sent_data["recipient_key"] == "recipient_key_abc"
        assert sent_data["body"] == {"text": "hi"}
        assert "request_id" in sent_data

    @pytest.mark.asyncio
    async def test_send_not_connected(self):
        """ws is None -- should raise ConnectionError."""
        keypair = Keypair.generate()
        ws = WSClient("http://localhost:8080", keypair)
        # ws._ws is None by default

        with pytest.raises(ConnectionError, match="Not connected"):
            await ws.send_message("recipient_key", {"text": "hi"})


class TestReconnectingWSClient:
    @pytest.mark.asyncio
    async def test_reconnects_on_disconnect(self):
        """First connect succeeds then read_loop raises ConnectionClosed, verify reconnection."""
        keypair = Keypair.generate()
        ws = ReconnectingWSClient(
            "http://localhost:8080", keypair,
            max_retries=3, base_delay=0.01, max_delay=0.05,
        )

        connect_count = 0

        async def mock_connect():
            nonlocal connect_count
            connect_count += 1

        call_count = 0

        async def mock_read_loop():
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                raise websockets.ConnectionClosed(None, None)
            else:
                # Second time, close cleanly to exit the loop
                ws._closed = True

        with patch.object(ws, "connect", side_effect=mock_connect):
            with patch.object(ws, "_read_loop", side_effect=mock_read_loop):
                await ws.run()

        assert connect_count >= 2

    @pytest.mark.asyncio
    async def test_max_retries(self):
        """connect always fails -- should stop after max_retries."""
        keypair = Keypair.generate()
        max_retries = 3
        ws = ReconnectingWSClient(
            "http://localhost:8080", keypair,
            max_retries=max_retries, base_delay=0.01, max_delay=0.05,
        )

        connect_count = 0

        async def mock_connect():
            nonlocal connect_count
            connect_count += 1
            raise ConnectionError("refused")

        with patch.object(ws, "connect", side_effect=mock_connect):
            await ws.run()

        assert connect_count == max_retries

    @pytest.mark.asyncio
    async def test_close_stops(self):
        """Call close() during run, verify loop ends."""
        keypair = Keypair.generate()
        ws = ReconnectingWSClient(
            "http://localhost:8080", keypair,
            max_retries=10, base_delay=0.01, max_delay=0.05,
        )

        async def mock_connect():
            pass

        async def mock_read_loop():
            # Simulate running for a bit, then close is called externally
            await asyncio.sleep(0.05)
            raise websockets.ConnectionClosed(None, None)

        async def close_after_delay():
            await asyncio.sleep(0.02)
            await ws.close()

        with patch.object(ws, "connect", side_effect=mock_connect):
            with patch.object(ws, "_read_loop", side_effect=mock_read_loop):
                # Run both concurrently -- close should cause run to exit
                await asyncio.gather(ws.run(), close_after_delay())

        assert ws._closed is True
