"""
Integration tests for CAAE with NVLink hardware.

Tests async communication patterns, queue depth measurement, and
circuit breaker behavior under actual GPU hardware constraints.

These tests are designed to run in staging environments with:
- Multi-GPU setups (H100, A100 with NVLink)
- Real PCIe transfers
- Network latency to CAAE service

Author: CAAE Research
Date: January 23, 2026
"""

import asyncio
import time
import pytest
from unittest.mock import Mock, AsyncMock, patch

from caae.vllm_integration import CAEBlockManagerWrapper
from caae.circuit_breaker import CircuitBreaker, CircuitState


class TestQueueDepthMeasurement:
    """Test suite for queue depth measurement in vLLM wrapper."""

    @pytest.fixture
    def mock_manager(self):
        """Create a mock block manager."""
        manager = Mock()
        manager.allocate_eviction_candidate = Mock(return_value=0)
        return manager

    @pytest.fixture
    def wrapper(self, mock_manager):
        """Create a wrapper with queue depth measurement enabled."""
        return CAEBlockManagerWrapper(
            original_manager=mock_manager,
            enable_queue_depth_measurement=True,
            timeout_ms=5.0,
        )

    def test_wrapper_initializes_queue_depth_tracking(self, wrapper):
        """Test: Wrapper initializes with queue depth tracking."""
        assert wrapper.enable_queue_depth_measurement is True
        assert wrapper._queue_depth_stats is not None
        assert wrapper._queue_depth_stats.total_measurements == 0

    def test_measure_queue_depth_records_samples(self, wrapper):
        """Test: Queue depth measurement records samples."""
        with patch.object(wrapper, "_get_pcie_queue_depth", return_value=1.5):
            # First measurement
            qd1 = wrapper._measure_queue_depth(context_tokens=25000)
            assert qd1 == 1.5

            # Wait for interval
            wrapper._last_measurement_time = 0

            # Second measurement
            qd2 = wrapper._measure_queue_depth(context_tokens=25000)
            assert qd2 == 1.5

            assert len(wrapper._queue_depth_samples) > 0

    def test_queue_depth_stats_calculations(self, wrapper):
        """Test: Queue depth stats are calculated correctly."""
        with patch.object(wrapper, "_get_pcie_queue_depth") as mock_qd:
            mock_qd.side_effect = [0.5, 1.0, 1.5, 2.0, 2.5]

            for i in range(5):
                wrapper._last_measurement_time = 0  # Force measurement
                wrapper._measure_queue_depth(context_tokens=25000)

            stats = wrapper.get_queue_depth_stats()

            assert stats.min_ms == 0.5
            assert stats.max_ms == 2.5
            assert stats.avg_ms == 1.5
            assert stats.p50_ms == 1.5
            assert stats.p99_ms == 2.5
            assert stats.total_measurements == 5

    def test_queue_depth_threshold_exceeded_tracking(self, wrapper):
        """Test: Threshold exceeded events are tracked."""
        wrapper.config.circuit_breaker_threshold_ms = 5.0

        with patch.object(wrapper, "_get_pcie_queue_depth") as mock_qd:
            # Below threshold
            mock_qd.return_value = 3.0
            wrapper._last_measurement_time = 0
            wrapper._measure_queue_depth(context_tokens=25000)

            stats = wrapper.get_queue_depth_stats()
            assert stats.threshold_exceeded_count == 0

            # Above threshold
            wrapper._last_measurement_time = 0
            mock_qd.return_value = 6.0
            wrapper._measure_queue_depth(context_tokens=25000)

            stats = wrapper.get_queue_depth_stats()
            assert stats.threshold_exceeded_count == 1

    def test_get_queue_depth_samples_filters_by_age(self, wrapper):
        """Test: get_queue_depth_samples respects max_age_seconds."""
        with patch.object(wrapper, "_get_pcie_queue_depth", return_value=1.0):
            # Add samples at different times
            for i in range(3):
                wrapper._last_measurement_time = 0
                wrapper._measure_queue_depth(context_tokens=25000)

            # Get all samples (recent)
            samples = wrapper.get_queue_depth_samples(max_age_seconds=60.0)
            assert len(samples) == 3

            # Get no samples (very old cutoff)
            samples = wrapper.get_queue_depth_samples(max_age_seconds=-1.0)
            assert len(samples) == 0

    def test_queue_depth_metrics_exported(self, wrapper, mock_manager):
        """Test: Queue depth metrics are included in get_metrics()."""
        with patch.object(wrapper, "_get_pcie_queue_depth", return_value=2.0):
            wrapper._last_measurement_time = 0
            wrapper._measure_queue_depth(context_tokens=25000)

            metrics = wrapper.get_metrics()

            assert "queue_depth_current_ms" in metrics
            assert "queue_depth_min_ms" in metrics
            assert "queue_depth_max_ms" in metrics
            assert "queue_depth_avg_ms" in metrics
            assert "queue_depth_p50_ms" in metrics
            assert "queue_depth_p95_ms" in metrics
            assert "queue_depth_p99_ms" in metrics
            assert "queue_depth_threshold_exceeded" in metrics


class TestAsyncCommunicationPatterns:
    """Test async communication patterns in vLLM wrapper."""

    @pytest.fixture
    def mock_manager(self):
        """Create a mock block manager."""
        manager = Mock()
        manager.allocate_eviction_candidate = Mock(return_value=0)
        return manager

    @pytest.fixture
    def wrapper(self, mock_manager):
        """Create a wrapper."""
        return CAEBlockManagerWrapper(
            original_manager=mock_manager,
            timeout_ms=100.0,  # Longer timeout for tests
        )

    @pytest.mark.asyncio
    async def test_async_query_caae_timeout_handling(self, wrapper):
        """Test: Async query handles timeouts gracefully."""
        import httpx

        with patch.object(wrapper, "_get_client") as mock_get_client:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout"))
            mock_get_client.return_value = mock_client

            result = await wrapper._query_caae(context_tokens=25000, queue_depth_ms=1.0)

            assert result is None
            assert wrapper.fallbacks_triggered == 1

    @pytest.mark.asyncio
    async def test_async_query_caae_error_handling(self, wrapper):
        """Test: Async query handles errors gracefully."""
        with patch.object(wrapper, "_get_client") as mock_get_client:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(side_effect=Exception("Connection error"))
            mock_get_client.return_value = mock_client

            result = await wrapper._query_caae(context_tokens=25000, queue_depth_ms=1.0)

            assert result is None
            assert wrapper.caae_errors == 1

    @pytest.mark.asyncio
    async def test_async_query_includes_queue_depth(self, wrapper):
        """Test: Async query includes queue depth in request."""
        expected_request_data = None

        async def capture_request(*args, **kwargs):
            nonlocal expected_request_data
            expected_request_data = kwargs.get("json", {})

            class MockResponse:
                async def raise_for_status(self):
                    pass

                def json(self):
                    return {
                        "action": "swap",
                        "predicted_swap_ms": 10.0,
                        "predicted_recompute_ms": 100.0,
                        "confidence": 0.95,
                        "reasoning": "Swap is faster",
                    }

            return MockResponse()

        with patch.object(wrapper, "_get_client") as mock_get_client:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(side_effect=capture_request)
            mock_get_client.return_value = mock_client

            await wrapper._query_caae(context_tokens=25000, queue_depth_ms=2.5)

            assert expected_request_data is not None
            assert expected_request_data["pcie_queue_depth_ms"] == 2.5
            assert expected_request_data["context_size_tokens"] == 25000

    @pytest.mark.asyncio
    async def test_concurrent_queries_to_caae(self, wrapper):
        """Test: Multiple concurrent queries to CAAE work correctly."""

        async def mock_query_success(*args, **kwargs):
            await asyncio.sleep(0.01)  # Simulate latency

            class MockResponse:
                async def raise_for_status(self):
                    pass

                def json(self):
                    return {
                        "action": "swap",
                        "predicted_swap_ms": 10.0,
                        "predicted_recompute_ms": 100.0,
                        "confidence": 0.95,
                        "reasoning": "Test response",
                    }

            return MockResponse()

        with patch.object(wrapper, "_get_client") as mock_get_client:
            mock_client = AsyncMock()
            mock_client.post = AsyncMock(side_effect=mock_query_success)
            mock_get_client.return_value = mock_client

            # Fire multiple concurrent queries
            tasks = [
                wrapper._query_caae(context_tokens=i * 1000, queue_depth_ms=0.5)
                for i in range(5)
            ]

            results = await asyncio.gather(*tasks)

            assert len(results) == 5
            assert all(r is not None for r in results)
            assert all(r["action"] == "swap" for r in results)


class TestCircuitBreakerWithQueueDepth:
    """Test circuit breaker interaction with queue depth measurement."""

    def test_circuit_breaker_activation_on_high_queue_depth(self):
        """Test: Circuit breaker activates when queue depth threshold exceeded."""
        cb = CircuitBreaker(threshold_ms=5.0, cooldown_ms=100.0)

        assert cb.state == CircuitState.CLOSED

        # Trigger with low queue depth
        result = cb.check(queue_depth_ms=2.0)
        assert result is False

        # Trigger with high queue depth multiple times
        for _ in range(2):
            result = cb.check(queue_depth_ms=6.0)

        # Should be open after threshold exceeded
        assert cb.is_active is True

    def test_circuit_breaker_recovery_with_queue_depth_improvement(self):
        """Test: Circuit breaker recovers when queue depth normalizes."""
        cb = CircuitBreaker(threshold_ms=5.0, cooldown_ms=10.0)

        # Trip the circuit
        for _ in range(2):
            cb.check(queue_depth_ms=6.0)

        assert cb.is_active is True

        # Wait for cooldown
        time.sleep(0.02)

        # Try recovery with good queue depth
        for _ in range(3):
            cb.check(queue_depth_ms=2.0)

        # Should be closed again
        assert cb.is_active is False

    def test_circuit_breaker_stats_include_thresholds(self):
        """Test: Circuit breaker stats track threshold exceedances."""
        cb = CircuitBreaker(threshold_ms=5.0)

        # Multiple checks at different levels
        cb.check(queue_depth_ms=2.0)  # Below threshold
        cb.check(queue_depth_ms=2.0)  # Below threshold
        cb.check(queue_depth_ms=6.0)  # Above threshold

        stats = cb.get_stats()

        assert stats.total_checks == 3


class TestMetricsExport:
    """Test metrics export for fallbacks and circuit breaker state."""

    @pytest.fixture
    def wrapper(self):
        """Create a wrapper."""
        manager = Mock()
        manager.allocate_eviction_candidate = Mock(return_value=0)

        return CAEBlockManagerWrapper(
            original_manager=manager,
            enable_queue_depth_measurement=True,
        )

    def test_metrics_include_fallback_counts(self, wrapper):
        """Test: Metrics include fallback trigger counts."""
        wrapper.fallbacks_triggered = 5
        wrapper.decisions_made = 100

        metrics = wrapper.get_metrics()

        assert metrics["fallbacks_triggered"] == 5
        assert metrics["fallback_ratio"] == 0.05

    def test_metrics_include_circuit_breaker_threshold_info(self, wrapper):
        """Test: Metrics include queue depth vs threshold info."""
        with patch.object(wrapper, "_get_pcie_queue_depth", return_value=6.0):
            wrapper.config.circuit_breaker_threshold_ms = 5.0
            wrapper._last_measurement_time = 0
            wrapper._measure_queue_depth(context_tokens=25000)

            metrics = wrapper.get_metrics()

            assert metrics["queue_depth_threshold_exceeded"] >= 0
            assert (
                metrics["queue_depth_current_ms"]
                > wrapper.config.circuit_breaker_threshold_ms
            )


# ============================================================================
# STAGING ENVIRONMENT TESTS (requires actual hardware)
# ============================================================================


@pytest.mark.skipif(
    True,  # Set to False when running in staging with NVLink hardware
    reason="Requires staging environment with NVLink hardware",
)
class TestNVLinkHardwareIntegration:
    """Integration tests that require actual NVLink hardware."""

    @pytest.fixture
    def nvlink_wrapper(self):
        """Create wrapper configured for NVLink hardware."""
        # In staging, this would connect to real vLLM instance
        manager = Mock()
        manager.allocate_eviction_candidate = Mock(return_value=0)

        return CAEBlockManagerWrapper(
            original_manager=manager,
            bandwidth_gbps=128.0,  # NVLink bandwidth
            enable_queue_depth_measurement=True,
        )

    @pytest.mark.asyncio
    async def test_real_nvlink_queue_depth_measurement(self, nvlink_wrapper):
        """Test: Measure queue depth on real NVLink hardware."""
        # This would measure actual PCIe queue depth from nvidia-smi
        # or NVIDIA driver APIs

        for _ in range(10):
            await asyncio.sleep(0.1)
            nvlink_wrapper._last_measurement_time = 0
            qd = nvlink_wrapper._measure_queue_depth(context_tokens=50000)

            # Validate we get numeric measurements
            assert isinstance(qd, (int, float))
            assert qd >= 0

    @pytest.mark.asyncio
    async def test_real_caae_decision_latency(self, nvlink_wrapper):
        """Test: Measure latency of actual CAAE decisions in staging."""
        times = []

        for _ in range(20):
            start = time.perf_counter()
            result = await nvlink_wrapper._query_caae(
                context_tokens=25000, queue_depth_ms=1.0
            )
            latency_ms = (time.perf_counter() - start) * 1000

            if result is not None:
                times.append(latency_ms)

        if times:
            p99_latency = sorted(times)[int(len(times) * 0.99)]

            # In staging, decision latency should be reasonable
            # (depends on network distance to CAAE service)
            assert p99_latency < 500.0, f"P99 latency {p99_latency}ms exceeds budget"

    @pytest.mark.asyncio
    async def test_circuit_breaker_prevents_pcie_saturation(self, nvlink_wrapper):
        """Test: Circuit breaker prevents excessive PCIe transfers."""
        # Run sustained load and verify circuit breaker activates
        # when queue depth exceeds threshold

        high_qd_count = 0

        for i in range(100):
            await asyncio.sleep(0.01)
            nvlink_wrapper._last_measurement_time = 0
            qd = nvlink_wrapper._measure_queue_depth(context_tokens=100000)

            if qd > nvlink_wrapper.config.circuit_breaker_threshold_ms:
                high_qd_count += 1

        # In normal operation, shouldn't exceed threshold constantly
        # but in saturation, circuit breaker should activate


if __name__ == "__main__":
    pytest.main([__file__, "-v", "-s"])
