# Copyright 2026 Ant Group Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for HttpCommunicator, CommConfig, CommStats, and timeout exceptions.

These tests use mocks to avoid actual HTTP communication, focusing on
the communicator's internal logic.
"""

import threading
import time
from unittest.mock import patch

import pytest

from mplang.backends.simp_worker.base import (
    RequestStatus,
    SendRequest,
    wait_all,
)
from mplang.backends.simp_worker.http import (
    CommConfig,
    CommStats,
    HttpCommunicator,
    RecvTimeoutError,
    SendTimeoutError,
)

# ---------------------------------------------------------------------------
# CommConfig Tests
# ---------------------------------------------------------------------------


class TestCommConfig:
    """Tests for CommConfig dataclass."""

    def test_default_values(self):
        """Test default configuration values."""
        config = CommConfig()
        assert config.default_recv_timeout == 600.0  # 10 minutes
        assert config.http_timeout == 60.0

    def test_custom_values(self):
        """Test custom configuration values."""
        config = CommConfig(
            default_recv_timeout=10.0,
            http_timeout=5.0,
        )
        assert config.default_recv_timeout == 10.0
        assert config.http_timeout == 5.0


# ---------------------------------------------------------------------------
# CommStats Tests
# ---------------------------------------------------------------------------


class TestCommStats:
    """Tests for CommStats dataclass."""

    def test_initial_values(self):
        """Test initial statistics are zero."""
        stats = CommStats()
        assert stats.messages_sent == 0
        assert stats.messages_received == 0
        assert stats.bytes_sent == 0
        assert stats.bytes_received == 0
        assert stats.send_errors == 0
        assert stats.recv_timeouts == 0
        assert stats.total_send_time_ms == 0.0
        assert stats.total_recv_wait_time_ms == 0.0

    def test_record_send(self):
        """Test recording send operations."""
        stats = CommStats()
        stats.record_send(1000, 50.0)
        stats.record_send(2000, 100.0)

        assert stats.messages_sent == 2
        assert stats.bytes_sent == 3000
        assert stats.total_send_time_ms == 150.0

    def test_record_recv(self):
        """Test recording receive operations."""
        stats = CommStats()
        stats.record_recv(500, 25.0)
        stats.record_recv(1500, 75.0)

        assert stats.messages_received == 2
        assert stats.bytes_received == 2000
        assert stats.total_recv_wait_time_ms == 100.0

    def test_record_send_error(self):
        """Test recording send errors."""
        stats = CommStats()
        stats.record_send_error()
        stats.record_send_error()

        assert stats.send_errors == 2

    def test_record_recv_timeout(self):
        """Test recording receive timeouts."""
        stats = CommStats()
        stats.record_recv_timeout()
        stats.record_recv_timeout()
        stats.record_recv_timeout()

        assert stats.recv_timeouts == 3

    def test_avg_send_latency_ms(self):
        """Test average send latency calculation."""
        stats = CommStats()
        assert stats.avg_send_latency_ms == 0.0  # No sends yet

        stats.record_send(100, 20.0)
        stats.record_send(100, 40.0)
        stats.record_send(100, 60.0)

        assert stats.avg_send_latency_ms == 40.0

    def test_avg_recv_wait_time_ms(self):
        """Test average receive wait time calculation."""
        stats = CommStats()
        assert stats.avg_recv_wait_time_ms == 0.0  # No receives yet

        stats.record_recv(100, 10.0)
        stats.record_recv(100, 30.0)

        assert stats.avg_recv_wait_time_ms == 20.0

    def test_snapshot(self):
        """Test snapshot returns correct dict."""
        stats = CommStats()
        stats.record_send(1000, 50.0)
        stats.record_recv(500, 25.0)
        stats.record_send_error()
        stats.record_recv_timeout()

        snapshot = stats.snapshot()

        assert snapshot["messages_sent"] == 1
        assert snapshot["messages_received"] == 1
        assert snapshot["bytes_sent"] == 1000
        assert snapshot["bytes_received"] == 500
        assert snapshot["send_errors"] == 1
        assert snapshot["recv_timeouts"] == 1
        assert snapshot["avg_send_latency_ms"] == 50.0
        assert snapshot["avg_recv_wait_time_ms"] == 25.0

    def test_reset(self):
        """Test reset clears all statistics."""
        stats = CommStats()
        stats.record_send(1000, 50.0)
        stats.record_recv(500, 25.0)
        stats.record_send_error()
        stats.record_recv_timeout()

        stats.reset()

        assert stats.messages_sent == 0
        assert stats.messages_received == 0
        assert stats.bytes_sent == 0
        assert stats.bytes_received == 0
        assert stats.send_errors == 0
        assert stats.recv_timeouts == 0
        assert stats.total_send_time_ms == 0.0
        assert stats.total_recv_wait_time_ms == 0.0

    def test_thread_safety(self):
        """Test that stats operations are thread-safe."""
        stats = CommStats()
        num_threads = 10
        ops_per_thread = 100

        def record_operations():
            for _ in range(ops_per_thread):
                stats.record_send(100, 1.0)
                stats.record_recv(50, 0.5)

        threads = [
            threading.Thread(target=record_operations) for _ in range(num_threads)
        ]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        expected_messages = num_threads * ops_per_thread
        assert stats.messages_sent == expected_messages
        assert stats.messages_received == expected_messages


# ---------------------------------------------------------------------------
# Exception Tests
# ---------------------------------------------------------------------------


class TestRecvTimeoutError:
    """Tests for RecvTimeoutError exception."""

    def test_attributes(self):
        """Test exception attributes are set correctly."""
        exc = RecvTimeoutError(frm=1, key="test_key", timeout=5.0)

        assert exc.frm == 1
        assert exc.key == "test_key"
        assert exc.timeout == 5.0

    def test_message(self):
        """Test exception message format."""
        exc = RecvTimeoutError(frm=2, key="data_key", timeout=10.0)

        assert "10.0s" in str(exc)
        assert "rank 2" in str(exc)
        assert "data_key" in str(exc)

    def test_is_timeout_error(self):
        """Test that RecvTimeoutError is a TimeoutError subclass."""
        exc = RecvTimeoutError(frm=0, key="k", timeout=1.0)
        assert isinstance(exc, TimeoutError)


class TestSendTimeoutError:
    """Tests for SendTimeoutError exception."""

    def test_attributes(self):
        """Test exception attributes are set correctly."""
        exc = SendTimeoutError(to=2, key="test_key", timeout=30.0)

        assert exc.to == 2
        assert exc.key == "test_key"
        assert exc.timeout == 30.0

    def test_message(self):
        """Test exception message format."""
        exc = SendTimeoutError(to=3, key="msg_key", timeout=15.0)

        assert "15.0s" in str(exc)
        assert "rank 3" in str(exc)
        assert "msg_key" in str(exc)

    def test_is_timeout_error(self):
        """Test that SendTimeoutError is a TimeoutError subclass."""
        exc = SendTimeoutError(to=0, key="k", timeout=1.0)
        assert isinstance(exc, TimeoutError)


# ---------------------------------------------------------------------------
# HttpCommunicator Tests (with mocks)
# ---------------------------------------------------------------------------


class TestHttpCommunicatorInit:
    """Tests for HttpCommunicator initialization."""

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_default_config(self, mock_client):
        """Test communicator with default config."""
        comm = HttpCommunicator(
            rank=0,
            world_size=3,
            endpoints=[
                "http://localhost:8000",
                "http://localhost:8001",
                "http://localhost:8002",
            ],
        )

        assert comm.rank == 0
        assert comm.world_size == 3
        assert comm.config.default_recv_timeout == 600.0  # 10 minutes
        assert comm.config.http_timeout == 60.0
        assert isinstance(comm.stats, CommStats)

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_custom_config(self, mock_client):
        """Test communicator with custom config."""
        config = CommConfig(default_recv_timeout=10.0, http_timeout=5.0)
        comm = HttpCommunicator(
            rank=1,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
            config=config,
        )

        assert comm.config.default_recv_timeout == 10.0
        assert comm.config.http_timeout == 5.0

        comm.shutdown()


class TestHttpCommunicatorRecvTimeout:
    """Tests for recv() timeout behavior."""

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_recv_timeout_raises_error(self, mock_client):
        """Test that recv raises RecvTimeoutError when timeout expires."""
        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        # recv with a short timeout should raise RecvTimeoutError
        with pytest.raises(RecvTimeoutError) as exc_info:
            comm.recv(frm=1, key="nonexistent", timeout=0.1)

        assert exc_info.value.frm == 1
        assert exc_info.value.key == "nonexistent"
        assert exc_info.value.timeout == 0.1

        # Stats should record the timeout
        assert comm.stats.recv_timeouts == 1

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_recv_uses_config_default_timeout(self, mock_client):
        """Test that recv uses config.default_recv_timeout when timeout not specified."""
        config = CommConfig(default_recv_timeout=0.1)
        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
            config=config,
        )

        # Should use default_recv_timeout from config
        with pytest.raises(RecvTimeoutError) as exc_info:
            comm.recv(frm=1, key="test")

        assert exc_info.value.timeout == 0.1

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_recv_explicit_timeout_overrides_config(self, mock_client):
        """Test that explicit timeout parameter overrides config default."""
        config = CommConfig(default_recv_timeout=10.0)
        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
            config=config,
        )

        # Explicit timeout should override config
        with pytest.raises(RecvTimeoutError) as exc_info:
            comm.recv(frm=1, key="test", timeout=0.05)

        assert exc_info.value.timeout == 0.05

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_recv_no_timeout_waits_for_data(self, mock_client):
        """Test that recv without timeout waits until data arrives."""
        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        test_data = {"key": "value"}

        def delayed_deliver():
            time.sleep(0.1)
            comm.on_receive(from_rank=1, key="delayed", data=test_data)

        t = threading.Thread(target=delayed_deliver)
        t.start()

        # Should block until data arrives
        result = comm.recv(frm=1, key="delayed")
        t.join()

        assert result == test_data
        assert comm.stats.messages_received == 1

        comm.shutdown()


class TestHttpCommunicatorOnReceive:
    """Tests for on_receive() method."""

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_on_receive_stores_data(self, mock_client):
        """Test that on_receive stores data in mailbox."""
        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        test_data = [1, 2, 3]
        comm.on_receive(from_rank=1, key="test_key", data=test_data)

        # Data should be retrievable via recv
        result = comm.recv(frm=1, key="test_key", timeout=1.0)
        assert result == test_data

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_on_receive_mailbox_overflow(self, mock_client):
        """Test that duplicate key raises RuntimeError."""
        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        comm.on_receive(from_rank=1, key="dup_key", data="first")

        with pytest.raises(RuntimeError, match="Mailbox overflow"):
            comm.on_receive(from_rank=1, key="dup_key", data="second")

        comm.shutdown()


class TestHttpCommunicatorStats:
    """Tests for statistics tracking in HttpCommunicator."""

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_recv_records_stats(self, mock_client):
        """Test that successful recv records statistics."""
        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        # Pre-populate mailbox
        comm.on_receive(from_rank=1, key="k1", data=b"hello")
        comm.on_receive(from_rank=1, key="k2", data="world")

        comm.recv(frm=1, key="k1", timeout=1.0)
        comm.recv(frm=1, key="k2", timeout=1.0)

        assert comm.stats.messages_received == 2
        # bytes data should record size
        assert comm.stats.bytes_received == 5  # len(b"hello")

        comm.shutdown()


class TestHttpCommunicatorSendSync:
    """Tests for send_sync() behavior."""

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_send_sync_success_records_stats(self, mock_client):
        """Test that send_sync succeeds and records stats."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        comm.send_sync(to=1, key="k", data=b"hi")

        assert comm.stats.messages_sent == 1
        assert comm.stats.send_errors == 0
        assert comm.stats.bytes_sent == 4  # len(base64.b64encode(b"hi"))

        _, kwargs = mock_client.return_value.put.call_args
        assert kwargs["json"]["is_raw_bytes"] is True
        # send_sync with no explicit timeout uses config.http_timeout (default 60.0)
        assert kwargs["timeout"] == 60.0

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_send_sync_timeout_raises_error(self, mock_client):
        """Test that send_sync raises SendTimeoutError on httpx timeout."""
        import httpx as _httpx

        # Make httpx.Client.put raise a TimeoutException
        mock_client.return_value.put.side_effect = _httpx.ReadTimeout("timed out")

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        with pytest.raises(SendTimeoutError) as exc_info:
            comm.send_sync(to=1, key="k", data="payload", timeout=0.01)

        assert exc_info.value.to == 1
        assert exc_info.value.key == "k"
        assert exc_info.value.timeout == 0.01
        # Stats should record exactly one error (no double-count)
        assert comm.stats.send_errors == 1
        assert comm.stats.messages_sent == 0

        comm.shutdown()


class TestHttpCommunicatorSend:
    """Tests for send() returning SendRequest (non-blocking send)."""

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_send_returns_send_request(self, mock_client):
        """Test that send() returns a SendRequest object."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        req = comm.send(to=1, key="test_key", data=b"payload")

        assert isinstance(req, SendRequest)
        assert req.to == 1
        assert req.key == "test_key"

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_send_request_wait(self, mock_client):
        """Test that SendRequest.wait() blocks until send completes."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        req = comm.send(to=1, key="k", data="data")
        result = req.wait()  # Should complete successfully

        assert result is None  # send returns None
        assert req.status == RequestStatus.COMPLETED
        assert comm.stats.messages_sent == 1

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_send_request_test_method(self, mock_client):
        """Test that SendRequest.test() returns completion status."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        req = comm.send(to=1, key="k", data="data")

        # Wait briefly for background thread to complete
        time.sleep(0.1)

        done, result = req.test()
        assert done is True
        assert result is None

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_multiple_sends_with_wait_all(self, mock_client):
        """Test multiple send() calls with wait_all()."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        reqs = [
            comm.send(to=1, key="k1", data="data1"),
            comm.send(to=1, key="k2", data="data2"),
            comm.send(to=1, key="k3", data="data3"),
        ]

        results = wait_all(reqs)

        assert len(results) == 3
        assert all(r is None for r in results)  # send returns None
        assert comm.stats.messages_sent == 3

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_isend_alias(self, mock_client):
        """Test that isend is an alias for send."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        # isend should be the same underlying function as send
        assert HttpCommunicator.isend is HttpCommunicator.send

        req = comm.isend(to=1, key="k", data="payload")
        assert isinstance(req, SendRequest)
        req.wait()

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_send_error_propagates_to_request(self, mock_client):
        """Test that send errors are captured in SendRequest."""
        import httpx as _httpx

        mock_client.return_value.put.side_effect = _httpx.ConnectError(
            "connection failed"
        )

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        req = comm.send(to=1, key="k", data="data")

        with pytest.raises(RuntimeError, match="Failed to send"):
            req.wait()

        assert req.status == RequestStatus.ERROR
        assert comm.stats.send_errors == 1

        comm.shutdown()


class TestPendingSendsPruning:
    """Tests for automatic pruning of completed futures from _pending_sends."""

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_completed_send_is_pruned(self, mock_client):
        """Test that a completed send future is removed from _pending_sends."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        req = comm.send(to=1, key="k", data="data")
        req.wait()  # Ensure the future completes

        # Give the done callback a moment to fire
        time.sleep(0.05)

        with comm._pending_sends_lock:
            assert len(comm._pending_sends) == 0

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_multiple_sends_all_pruned(self, mock_client):
        """Test that multiple completed sends are all pruned."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        reqs = [comm.send(to=1, key=f"k{i}", data=f"d{i}") for i in range(10)]
        wait_all(reqs)

        # Give done callbacks time to fire
        time.sleep(0.05)

        with comm._pending_sends_lock:
            assert len(comm._pending_sends) == 0

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_failed_send_is_also_pruned(self, mock_client):
        """Test that a failed send future is also removed from _pending_sends."""
        import httpx as _httpx

        mock_client.return_value.put.side_effect = _httpx.ConnectError(
            "connection failed"
        )

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        req = comm.send(to=1, key="k", data="data")

        with pytest.raises(RuntimeError, match="Failed to send"):
            req.wait()

        # Give the done callback a moment to fire
        time.sleep(0.05)

        with comm._pending_sends_lock:
            assert len(comm._pending_sends) == 0

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_pending_sends_is_set(self, mock_client):
        """Test that _pending_sends is a set (not a list) for O(1) removal."""
        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        assert isinstance(comm._pending_sends, set)

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_wait_pending_sends_without_pruning_still_works(self, mock_client):
        """Test that wait_pending_sends works correctly with the pruning mechanism."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        for i in range(5):
            comm.send(to=1, key=f"k{i}", data=f"d{i}")

        # wait_pending_sends should still work and not raise
        comm.wait_pending_sends()

        # After everything completes, set should be empty
        time.sleep(0.05)
        with comm._pending_sends_lock:
            assert len(comm._pending_sends) == 0

        comm.shutdown()

    @patch("mplang.backends.simp_worker.http.httpx.Client")
    def test_concurrent_send_and_pruning(self, mock_client):
        """Test that concurrent sends and pruning don't cause race conditions."""
        mock_response = mock_client.return_value.put.return_value
        mock_response.raise_for_status.return_value = None

        comm = HttpCommunicator(
            rank=0,
            world_size=2,
            endpoints=["http://localhost:8000", "http://localhost:8001"],
        )

        # Fire many sends concurrently
        reqs = [comm.send(to=1, key=f"k{i}", data=f"d{i}") for i in range(50)]
        wait_all(reqs)

        # Give done callbacks time to fire
        time.sleep(0.1)

        with comm._pending_sends_lock:
            assert len(comm._pending_sends) == 0

        assert comm.stats.messages_sent == 50

        comm.shutdown()
