"""
Tests for RunnerClient heartbeat/reconnect loop under real async conditions.

Uses respx to mock HTTP heartbeat responses and a mock WebSocket server
to exercise the full reconnection and backoff logic.
"""

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

import httpx
import pytest
import respx

from skop_runner.client import RunnerClient
from skop_runner.config import RunnerConfig
from skop_runner.output_buffer import OutputBuffer


@pytest.fixture
def mock_config(tmp_path):
    """Create a mock config for testing."""
    return RunnerConfig(
        device_token="test-token-123",
        api_url="http://localhost:3000",
        ws_url="ws://localhost:8080/runner",
        log_level="DEBUG",
        config_dir=tmp_path,
    )


@pytest.fixture
def mock_message_handler():
    """Create a mock message handler."""
    handler = MagicMock()
    handler.handle_message = AsyncMock(return_value={"jsonrpc": "2.0", "id": 1, "result": {}})
    return handler


class MockWebSocket:
    """Mock WebSocket that simulates connection behavior."""

    def __init__(self, messages=None, close_after=None):
        self.messages = messages or []
        self.close_after = close_after
        self.sent = []
        self.closed = False
        self._message_index = 0

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

    async def close(self):
        self.closed = True

    def __aiter__(self):
        return self

    async def __anext__(self):
        if self.closed:
            raise StopAsyncIteration

        if self.close_after is not None and self._message_index >= self.close_after:
            from websockets.exceptions import ConnectionClosed

            raise ConnectionClosed(None, None)

        if self._message_index < len(self.messages):
            msg = self.messages[self._message_index]
            self._message_index += 1
            return msg

        # No more messages, wait indefinitely (will be cancelled)
        await asyncio.sleep(10)
        raise StopAsyncIteration


@pytest.mark.asyncio
async def test_heartbeat_success(mock_config, mock_message_handler):
    """Test heartbeat sends successfully and logs debug message."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())
    client.running = True

    heartbeat_count = 0

    with respx.mock:
        # Mock successful heartbeat
        route = respx.post("http://localhost:3000/api/runner/heartbeat").mock(
            return_value=httpx.Response(200, json={"success": True})
        )

        async def stop_after_heartbeat():
            nonlocal heartbeat_count
            # Give time for one heartbeat cycle
            await asyncio.sleep(0.1)
            heartbeat_count = route.call_count
            client.running = False

        # Override sleep to speed up test
        original_sleep = asyncio.sleep

        async def fast_sleep(seconds):
            # Convert 30s heartbeat interval to 0.05s for testing
            # Heartbeat interval is 10 seconds in client.py
            await original_sleep(0.05 if seconds == 10 else seconds)

        with patch("asyncio.sleep", fast_sleep):
            await asyncio.gather(
                client.heartbeat_loop(),
                stop_after_heartbeat(),
            )

        assert heartbeat_count >= 1
        assert route.called


@pytest.mark.asyncio
async def test_heartbeat_failure_recovery(mock_config, mock_message_handler):
    """Test heartbeat handles failures gracefully and continues."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())
    client.running = True

    call_results = []

    with respx.mock:
        # First call fails, second succeeds
        def side_effect(request):
            call_num = len(call_results)
            if call_num == 0:
                call_results.append("fail_500")
                return httpx.Response(500, json={"error": "Server error"})
            elif call_num == 1:
                call_results.append("fail_timeout")
                raise httpx.TimeoutException("Connection timeout")
            else:
                call_results.append("success")
                return httpx.Response(200, json={"success": True})

        respx.post("http://localhost:3000/api/runner/heartbeat").mock(side_effect=side_effect)

        async def stop_after_recovery():
            # Wait for enough heartbeat cycles (at least 4)
            await asyncio.sleep(0.3)
            client.running = False

        original_sleep = asyncio.sleep

        async def fast_sleep(seconds):
            # Heartbeat interval is 10 seconds in client.py
            await original_sleep(0.05 if seconds == 10 else seconds)

        with patch("asyncio.sleep", fast_sleep):
            await asyncio.gather(
                client.heartbeat_loop(),
                stop_after_recovery(),
            )

        # Should have tried multiple times: fail, fail, success
        assert len(call_results) >= 3
        assert "fail_500" in call_results
        assert "fail_timeout" in call_results
        assert "success" in call_results


@pytest.mark.asyncio
async def test_reconnect_with_exponential_backoff(mock_config, mock_message_handler):
    """Test run_with_reconnect uses exponential backoff on connection failures."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())

    connect_attempts = []
    sleep_durations = []

    original_sleep = asyncio.sleep

    async def tracking_sleep(seconds):
        sleep_durations.append(seconds)
        # Actually sleep a tiny bit to allow task switching
        await original_sleep(0.01)

    async def failing_connect():
        connect_attempts.append(len(connect_attempts) + 1)
        # Fail the first 3 attempts
        if len(connect_attempts) <= 3:
            return False
        # Stop after 4th attempt
        client.running = False
        return False

    client.connect = failing_connect

    with patch("random.uniform", return_value=0):  # deterministic jitter
        with patch("asyncio.sleep", tracking_sleep):
            await client.run_with_reconnect()

    # Should have attempted to connect multiple times
    assert len(connect_attempts) >= 3

    # Verify exponential backoff: 1, 2, 4, ... (jitter is 0)
    assert sleep_durations[0] == 1
    assert sleep_durations[1] == 2
    assert sleep_durations[2] == 4


@pytest.mark.asyncio
async def test_reconnect_max_backoff_capped(mock_config, mock_message_handler):
    """Test reconnect backoff is capped at max delay."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())
    client._reconnect_delay = 4  # Start near cap

    sleep_durations = []
    connect_attempts = 0

    original_sleep = asyncio.sleep

    async def tracking_sleep(seconds):
        sleep_durations.append(seconds)
        await original_sleep(0.01)

    async def failing_connect():
        nonlocal connect_attempts
        connect_attempts += 1
        if connect_attempts >= 4:
            client.running = False
        return False

    client.connect = failing_connect

    with patch("random.uniform", return_value=0):  # deterministic jitter
        with patch("asyncio.sleep", tracking_sleep):
            await client.run_with_reconnect()

    # Backoff should cap at 15 seconds (with jitter=0)
    for duration in sleep_durations:
        assert duration <= 15


@pytest.mark.asyncio
async def test_successful_connect_resets_backoff(mock_config, mock_message_handler):
    """Test successful connection resets the backoff delay."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())
    client._reconnect_delay = 32  # Start with high delay

    connect_count = 0
    on_connect_called = False

    def on_connect():
        nonlocal on_connect_called
        on_connect_called = True

    client.on_connect = on_connect

    async def successful_connect():
        nonlocal connect_count
        connect_count += 1
        # Simulate successful connection that immediately "disconnects"
        client._reconnect_delay = 1  # This happens in real connect()
        return True

    async def immediate_exit_receive():
        # Simulate connection closing immediately
        client.running = False

    async def no_op_heartbeat():
        pass

    client.connect = successful_connect
    client.receive_loop = immediate_exit_receive
    client.heartbeat_loop = no_op_heartbeat

    await client.run_with_reconnect()

    assert connect_count == 1
    assert on_connect_called
    assert client._reconnect_delay == 1  # Reset to initial


@pytest.mark.asyncio
async def test_graceful_shutdown_during_reconnect(mock_config, mock_message_handler):
    """Test graceful shutdown while waiting for reconnect."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())

    connect_attempts = 0

    async def failing_connect():
        nonlocal connect_attempts
        connect_attempts += 1
        return False

    client.connect = failing_connect

    async def shutdown_during_sleep():
        # Wait a bit then trigger shutdown
        await asyncio.sleep(0.05)
        await client.disconnect()

    # Run both concurrently
    await asyncio.gather(
        client.run_with_reconnect(),
        shutdown_during_sleep(),
    )

    # Should have stopped cleanly
    assert client.running is False
    assert connect_attempts >= 1


@pytest.mark.asyncio
async def test_full_lifecycle_with_heartbeat_and_messages(mock_config, mock_message_handler):
    """Test full lifecycle: connect, receive messages, heartbeat, disconnect."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())

    # Track events
    events = []

    # Mock WebSocket
    mock_ws = MockWebSocket(
        messages=[
            '{"jsonrpc": "2.0", "id": 1, "method": "file.list", "params": {}}',
            '{"jsonrpc": "2.0", "id": 2, "method": "kernel.start", "params": {}}',
        ],
        close_after=2,  # Close after 2 messages
    )

    async def mock_connect_fn(*args, **kwargs):
        events.append("ws_connected")
        return mock_ws

    with respx.mock:
        # Mock heartbeat endpoint
        heartbeat_route = respx.post("http://localhost:3000/api/runner/heartbeat").mock(
            return_value=httpx.Response(200, json={"success": True})
        )

        with patch("websockets.connect", mock_connect_fn):
            # Speed up heartbeat interval
            original_sleep = asyncio.sleep

            async def fast_sleep(seconds):
                # Heartbeat interval is 10 seconds in client.py
                if seconds == 10:
                    events.append("heartbeat_sleep")
                    await original_sleep(0.02)
                else:
                    await original_sleep(min(seconds, 0.02))

            with patch("asyncio.sleep", fast_sleep):

                async def stop_after_messages():
                    # Wait for messages to be processed
                    await asyncio.sleep(0.15)
                    await client.disconnect()

                await asyncio.gather(
                    client.run_with_reconnect(),
                    stop_after_messages(),
                )

        # Verify lifecycle
        assert "ws_connected" in events
        assert heartbeat_route.called
        assert mock_message_handler.handle_message.call_count >= 2
        assert len(mock_ws.sent) >= 2  # Responses sent back


@pytest.mark.asyncio
async def test_heartbeat_timeout_recovery(mock_config, mock_message_handler):
    """Test heartbeat recovers from timeout errors."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())
    client.running = True

    responses = []

    with respx.mock:
        call_count = 0

        def timeout_then_success(request):
            nonlocal call_count
            call_count += 1
            if call_count <= 2:
                responses.append("timeout")
                raise httpx.TimeoutException("Timeout")
            responses.append("success")
            return httpx.Response(200, json={"success": True})

        respx.post("http://localhost:3000/api/runner/heartbeat").mock(side_effect=timeout_then_success)

        async def stop_after_recovery():
            # Wait for enough heartbeat cycles (at least 4)
            await asyncio.sleep(0.3)
            client.running = False

        original_sleep = asyncio.sleep

        async def fast_sleep(seconds):
            # Heartbeat interval is 10 seconds in client.py
            await original_sleep(0.05 if seconds == 10 else seconds)

        with patch("asyncio.sleep", fast_sleep):
            await asyncio.gather(
                client.heartbeat_loop(),
                stop_after_recovery(),
            )

    # Should have recovered from timeouts
    assert responses.count("timeout") >= 2
    assert "success" in responses


@pytest.mark.asyncio
async def test_concurrent_heartbeat_and_receive(mock_config, mock_message_handler):
    """Test heartbeat and receive loops run concurrently without interference."""
    client = RunnerClient(mock_config, mock_message_handler, OutputBuffer())

    heartbeat_calls = []
    messages_received = []

    # Create mock WebSocket with messages
    mock_ws = MockWebSocket(
        messages=[
            '{"jsonrpc": "2.0", "id": 1, "method": "test.method", "params": {}}',
        ] * 5,
        close_after=5,
    )

    client.ws = mock_ws
    client.running = True

    with respx.mock:
        heartbeat_route = respx.post("http://localhost:3000/api/runner/heartbeat").mock(
            side_effect=lambda req: (
                heartbeat_calls.append(1),
                httpx.Response(200, json={"success": True}),
            )[1]
        )

        original_handle = mock_message_handler.handle_message

        async def tracking_handle(msg):
            messages_received.append(msg)
            return await original_handle(msg)

        mock_message_handler.handle_message = tracking_handle

        original_sleep = asyncio.sleep

        async def fast_sleep(seconds):
            # Heartbeat interval is 10 seconds in client.py
            await original_sleep(0.02 if seconds == 10 else seconds)

        with patch("asyncio.sleep", fast_sleep):

            async def stop_after_work():
                await asyncio.sleep(0.2)
                client.running = False
                mock_ws.closed = True

            await asyncio.gather(
                client.receive_loop(),
                client.heartbeat_loop(),
                stop_after_work(),
                return_exceptions=True,
            )

    # Both loops should have done work
    assert len(heartbeat_calls) >= 1
    assert len(messages_received) >= 1
