"""Tests for the StreamModule."""

from __future__ import annotations

import asyncio
import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from openhedge.modules.stream import StreamModule


FAKE_FEED_URL = "wss://api.openhedge.io/feed"
FAKE_API_KEY = "oh_test_key"


class TestStreamConnect:
    async def test_connect_success(self):
        stream = StreamModule(feed_url=FAKE_FEED_URL, api_key=FAKE_API_KEY)

        mock_ws = AsyncMock()
        mock_ws.recv = AsyncMock(return_value=json.dumps({"type": "connected"}))

        with patch("openhedge.modules.stream.websockets.connect", new_callable=AsyncMock, return_value=mock_ws):
            await stream.connect()

        assert stream._connected is True
        # Verify auth was sent as in-band message (not via headers)
        mock_ws.send.assert_called_once()
        auth_msg = json.loads(mock_ws.send.call_args[0][0])
        assert auth_msg["type"] == "auth"
        assert auth_msg["apiKey"] == FAKE_API_KEY

    async def test_subscribe_before_connect_raises(self):
        stream = StreamModule(feed_url=FAKE_FEED_URL, api_key=FAKE_API_KEY)
        with pytest.raises(RuntimeError, match="Not connected"):
            await stream.subscribe()


class TestStreamEventHandlers:
    def test_on_decorator_registers_handler(self):
        stream = StreamModule(feed_url=FAKE_FEED_URL, api_key=FAKE_API_KEY)

        @stream.on("vault_nav_update")
        async def handler(data):
            pass

        assert len(stream._handlers["vault_nav_update"]) == 1

    def test_multiple_handlers(self):
        stream = StreamModule(feed_url=FAKE_FEED_URL, api_key=FAKE_API_KEY)

        @stream.on("vault_nav_update")
        async def handler1(data):
            pass

        @stream.on("vault_nav_update")
        async def handler2(data):
            pass

        assert len(stream._handlers["vault_nav_update"]) == 2


class TestStreamDisconnect:
    async def test_disconnect(self):
        stream = StreamModule(feed_url=FAKE_FEED_URL, api_key=FAKE_API_KEY)
        mock_ws = AsyncMock()
        stream._ws = mock_ws
        stream._connected = True

        await stream.disconnect()

        mock_ws.close.assert_called_once()
        assert stream._connected is False
        assert stream._ws is None

    async def test_disconnect_when_not_connected(self):
        stream = StreamModule(feed_url=FAKE_FEED_URL, api_key=FAKE_API_KEY)
        await stream.disconnect()  # should not raise


class TestStreamSubscribe:
    async def test_subscribe_sends_message(self):
        stream = StreamModule(feed_url=FAKE_FEED_URL, api_key=FAKE_API_KEY)
        mock_ws = AsyncMock()
        stream._ws = mock_ws

        await stream.subscribe(vaults=["0x" + "11" * 20])

        mock_ws.send.assert_called_once()
        sent = json.loads(mock_ws.send.call_args[0][0])
        assert sent["type"] == "subscribe"
        assert len(sent["vaults"]) == 1
