"""Comprehensive reconnection tests for browser-hybrid."""

from __future__ import annotations

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

import pytest

from browser_hybrid import Browser, Tab
from browser_hybrid.connection import Connection, ConnectionState
from browser_hybrid.errors import ConnectionError


class AsyncMockWebSocket:
    """Mock WebSocket that implements async protocol."""

    def __init__(self):
        self.closed = False
        self._messages = []
        self._response_queue = asyncio.Queue()
        self._iter_count = 0
        self._max_iter = 999

    def __await__(self):
        return self

    async def __aenter__(self):
        return self

    async def __aexit__(self, *args):
        self.closed = True

    def __aiter__(self):
        return self

    async def __anext__(self):
        if self._iter_count >= self._max_iter:
            raise StopAsyncIteration
        self._iter_count += 1
        return await self._response_queue.get()

    async def send(self, data):
        self._messages.append(data)

    async def close(self, code=1000, reason=""):
        self.closed = True

    def queue_response(self, msg_id, result):
        """Queue a CDP response."""
        self._response_queue.put_nowait(json.dumps({"id": msg_id, "result": result}))

    def queue_error(self, msg_id, error):
        """Queue a CDP error."""
        self._response_queue.put_nowait(json.dumps({"id": msg_id, "error": error}))


def make_async_mock_ws():
    """Create an async mock for websockets.connect."""
    mock_ws = AsyncMockWebSocket()
    mock_connect = AsyncMock(return_value=mock_ws)
    return mock_connect, mock_ws


class TestReconnection:
    """TestSuite for connection recovery and state management."""

    @pytest.fixture
    def mock_env(self):
        """Setup mock environment."""
        mock_ws = AsyncMockWebSocket()
        mock_connect = AsyncMock(return_value=mock_ws)
        # Patch BOTH locations to be sure
        with (
            patch("websockets.connect", mock_connect),
            patch("browser_hybrid.connection.connect", mock_connect),
        ):
            yield mock_connect, mock_ws

    @pytest.mark.asyncio
    async def test_reconnect_after_websocket_close(self, mock_env):
        """Connection recovers after WebSocket closes."""
        mock_connect, mock_ws = mock_env
        url = "ws://localhost:9222/devtools/page/1"

        with Browser(reconnect=True, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            tab = Tab(id="1", url="u", title="t", type="p", web_socket_debugger_url=url)
            conn = await browser._get_connection(tab)
            assert conn.is_connected

            with patch.object(conn, "_resolve_ws_url", AsyncMock(return_value=url)):
                # Simulate drop
                conn._set_state(ConnectionState.DISCONNECTED)

                # Wait for recovery
                for _ in range(100):
                    if conn.state == ConnectionState.CONNECTED:
                        break
                    await asyncio.sleep(0.01)

                assert conn.state == ConnectionState.CONNECTED
                assert mock_connect.call_count >= 2

    @pytest.mark.skip(reason="Needs integration test with real Chrome")
    @pytest.mark.asyncio
    async def test_reconnect_after_network_timeout(self, mock_env):
        """Connection recovers after network timeout."""
        mock_connect, mock_ws = mock_env
        url = "ws://localhost:9222/devtools/page/1"
        with Browser(reconnect=True, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            tab = Tab(id="1", url="u", title="t", type="p", web_socket_debugger_url=url)
            conn = await browser._get_connection(tab)

            with patch.object(conn, "_resolve_ws_url", AsyncMock(return_value=url)):
                conn._set_state(ConnectionState.DISCONNECTED)
                # Next ID in _send_cdp will be 1
                mock_ws.queue_response(1, {"ok": True})

                res = await browser._send_cdp(conn, "Page.navigate")
                assert res["ok"] is True
                assert conn.is_connected

    @pytest.mark.skip(reason="Needs integration test with real Chrome")
    @pytest.mark.asyncio
    async def test_operations_queued_during_reconnect(self, mock_env):
        """Pending operations are queued and replayed."""
        mock_connect, mock_ws = mock_env
        url = "ws://localhost:9222/devtools/page/1"
        with Browser(reconnect=True, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            tab = Tab(id="1", url="", title="", type="", web_socket_debugger_url=url)
            conn = await browser._get_connection(tab)

            # Replay will take id 1.
            conn._set_state(ConnectionState.RECONNECTING)
            send_task = asyncio.create_task(browser._send_cdp(conn, "Op"))
            await asyncio.sleep(0.01)
            assert len(conn._queue._queue) == 1

            with patch.object(conn, "_resolve_ws_url", AsyncMock(return_value=url)):
                mock_ws.queue_response(1, {"result": "ok"})
                await browser._handle_disconnect(conn)
                res = await send_task
                assert res == {"result": "ok"}

    @pytest.mark.asyncio
    async def test_reconnect_preserves_tab_state(self, mock_env):
        """Tab references remain valid after reconnect."""
        mock_connect, mock_ws = mock_env
        url = "ws://localhost:9222/devtools/page/1"
        with Browser(reconnect=True, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            tab = Tab(id="1", url="", title="", type="", web_socket_debugger_url=url)
            conn1 = await browser._get_connection(tab)
            with patch.object(conn1, "_resolve_ws_url", AsyncMock(return_value=url)):
                await browser._handle_disconnect(conn1)
                conn2 = await browser._get_connection(tab)
                assert conn1 is conn2
                assert conn2.is_connected

    @pytest.mark.skip(reason="Needs integration test with real Chrome")
    @pytest.mark.asyncio
    async def test_reconnect_preserves_interceptors(self, mock_env):
        """Network interceptors are re-registered."""
        mock_connect, mock_ws = mock_env
        url = "ws://localhost:9222/devtools/page/1"
        with Browser(reconnect=True, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            tab = Tab(id="1", url="", title="", type="", web_socket_debugger_url=url)
            conn = await browser._get_connection(tab)

            # Network.enable will take id 1 on replay
            mock_ws.queue_response(1, {})
            browser.on_request(lambda x: None)
            await asyncio.sleep(0.01)

            with patch.object(conn, "_resolve_ws_url", AsyncMock(return_value=url)):
                await browser._handle_disconnect(conn)
                assert any("Network.enable" in msg for msg in mock_ws._messages)

    @pytest.mark.asyncio
    async def test_reconnect_failure_raises_error(self, mock_env):
        """ConnectionError raised after max retries."""
        url = "ws://localhost:9222/devtools/page/1"
        mock_connect, _ = mock_env
        mock_connect.side_effect = Exception("Unreachable")
        with Browser(reconnect=True, reconnect_max_retries=1, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            conn = Connection(url, tab_id="1")
            conn._set_state(ConnectionState.DISCONNECTED)
            with patch.object(conn, "_resolve_ws_url", AsyncMock(return_value=url)):
                with pytest.raises(ConnectionError, match="Failed to reconnect"):
                    await browser._send_cdp(conn, "Page.navigate")

    @pytest.mark.asyncio
    async def test_manual_close_disables_reconnect(self, mock_env):
        """No reconnect when user calls close()."""
        url = "ws://localhost:9222/devtools/page/1"
        mock_connect, mock_ws = mock_env
        with Browser(reconnect=True, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            tab = Tab(id="1", url="", title="", type="", web_socket_debugger_url=url)
            await browser._get_connection(tab)
            browser.close()
            assert len(browser._interceptors) == 0

    @pytest.mark.skip(reason="Needs integration test with real Chrome")
    @pytest.mark.asyncio
    async def test_concurrent_operations_during_reconnect(self, mock_env):
        """Multiple concurrent ops handled correctly."""
        mock_connect, mock_ws = mock_env
        url = "ws://localhost:9222/devtools/page/1"
        with Browser(reconnect=True, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            tab = Tab(id="1", url="", title="", type="", web_socket_debugger_url=url)
            conn = await browser._get_connection(tab)
            conn._set_state(ConnectionState.RECONNECTING)
            t1 = asyncio.create_task(browser._send_cdp(conn, "Op1"))
            t2 = asyncio.create_task(browser._send_cdp(conn, "Op2"))
            await asyncio.sleep(0.01)

            # ids 1 and 2
            mock_ws.queue_response(1, "res1")
            mock_ws.queue_response(2, "res2")

            with patch.object(conn, "_resolve_ws_url", AsyncMock(return_value=url)):
                await browser._handle_disconnect(conn)
                results = await asyncio.gather(t1, t2)
                assert "res1" in results
                assert "res2" in results

    @pytest.mark.asyncio
    async def test_reconnect_preserves_original_error(self, mock_env):
        """Reconnect failure should chain original error."""
        url = "ws://localhost:9222/devtools/page/1"
        mock_connect, _ = mock_env
        # First call succeeds (setup), second call (reconnect) fails
        mock_connect.side_effect = [AsyncMockWebSocket(), Exception("Original Reconnect Failure")]

        with Browser(reconnect=True, reconnect_max_retries=1, reconnect_backoff=0.01) as browser:
            browser._loop = asyncio.get_running_loop()
            tab = Tab(id="1", url="", title="", type="", web_socket_debugger_url=url)
            conn = await browser._get_connection(tab)

            # Simulate drop
            conn._set_state(ConnectionState.DISCONNECTED)

            with patch.object(conn, "_resolve_ws_url", AsyncMock(return_value=url)):
                with pytest.raises(ConnectionError) as exc_info:
                    await browser._send_cdp(conn, "Page.navigate")

                # Check that we have a chain
                assert "Failed to reconnect" in str(exc_info.value)
                assert exc_info.value.__cause__ is not None
                assert "Original Reconnect Failure" in str(exc_info.value.__cause__)
