"""Unit tests for intelligent batching."""

import pytest
import asyncio
from unittest.mock import AsyncMock
from azure_discovery.utils.batching import BatchProcessor
from azure_discovery.utils.rate_control import AdaptiveRateController, RateLimitConfig


class TestBatchProcessor:
    """Tests for batch processor."""
    
    @pytest.mark.asyncio
    async def test_processes_all_items(self):
        """Test that all items are processed."""
        processor = BatchProcessor(
            initial_batch_size=10,
            max_concurrent=2,
        )
        
        items = list(range(50))
        processed = []
        
        async def mock_processor(batch):
            processed.extend(batch)
            return len(batch)
        
        results = []
        async for result in processor.process_batches(items, mock_processor):
            results.append(result)
        
        assert len(processed) == 50
        assert set(processed) == set(items)
    
    @pytest.mark.asyncio
    async def test_increases_batch_size_on_success(self):
        """Test batch size increases after successful processing."""
        processor = BatchProcessor(
            initial_batch_size=10,
            max_batch_size=100,
            max_concurrent=1,
        )
        
        initial_size = processor.get_current_batch_size()
        
        items = list(range(20))
        
        async def mock_processor(batch):
            return len(batch)
        
        async for _ in processor.process_batches(items, mock_processor):
            pass
        
        final_size = processor.get_current_batch_size()
        assert final_size > initial_size
    
    @pytest.mark.asyncio
    async def test_decreases_batch_size_on_throttle(self):
        """Test batch size decreases when throttled."""
        processor = BatchProcessor(
            initial_batch_size=100,
            min_batch_size=10,
            max_concurrent=1,
        )

        initial_size = processor.get_current_batch_size()
        sizes_after_throttle: list[int] = []
        call_count = 0

        items = list(range(200))

        async def mock_processor(batch):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                # Simulate throttling on first batch
                raise Exception("429 Too Many Requests")
            # Record size after throttle-induced decrease
            sizes_after_throttle.append(processor.get_current_batch_size())
            return len(batch)

        try:
            async for _ in processor.process_batches(items, mock_processor):
                pass
        except Exception:
            pass  # Expected on first batch

        # The batch size should have decreased after the throttle event
        assert len(sizes_after_throttle) > 0
        assert sizes_after_throttle[0] < initial_size
    
    @pytest.mark.asyncio
    async def test_respects_max_concurrent(self):
        """Test concurrent batch limit is enforced."""
        max_concurrent = 3
        processor = BatchProcessor(
            initial_batch_size=10,
            max_concurrent=max_concurrent,
        )
        
        active_count = 0
        max_active = 0
        
        async def mock_processor(batch):
            nonlocal active_count, max_active
            active_count += 1
            max_active = max(max_active, active_count)
            await asyncio.sleep(0.1)  # Simulate work
            active_count -= 1
            return len(batch)
        
        items = list(range(100))
        
        async for _ in processor.process_batches(items, mock_processor):
            pass
        
        assert max_active <= max_concurrent
    
    @pytest.mark.asyncio
    async def test_integration_with_rate_controller(self):
        """Test batch processor integrates with rate controller."""
        rate_controller = AdaptiveRateController(
            RateLimitConfig(initial_rps=100.0)
        )
        
        processor = BatchProcessor(
            initial_batch_size=10,
            max_concurrent=2,
        )
        
        items = list(range(50))
        
        async def mock_processor(batch):
            return len(batch)
        
        results = []
        async for result in processor.process_batches(
            items, mock_processor, rate_controller
        ):
            results.append(result)
        
        assert len(results) > 0
        assert rate_controller.get_current_rps() > 0
