"""
Comprehensive tests for LLMLogger.

Tests cover:
- Basic functionality
- Thread safety
- Batching behavior
- Retry logic
- Graceful shutdown
- Queue overflow handling
"""

import json
import pytest
import threading
import time
from unittest.mock import Mock, patch, MagicMock
from queue import Queue

from burt_logger import LLMLogger


class TestLLMLoggerBasics:
    """Test basic logger functionality."""
    
    def test_initialization(self):
        """Test logger initializes with correct parameters."""
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=5,
            flush_interval=2.0,
        )
        
        assert logger.endpoint == "https://test.com/logs"
        assert logger.api_key == "test-key"
        assert logger.batch_size == 5
        assert logger.flush_interval == 2.0
        assert logger._worker_thread is not None
        assert logger._worker_thread.is_alive()
        
        logger.shutdown()
    
    def test_log_queues_entry(self):
        """Test that log() successfully queues entries."""
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
        )
        
        request = {"model": "gpt-3.5-turbo", "prompt": "test"}
        response = {"completion": "test response"}
        
        success = logger.log(request=request, response=response)
        
        assert success is True
        assert logger.stats["logs_queued"] == 1
        assert logger.log_queue.qsize() == 1
        
        logger.shutdown()
    
    def test_log_with_metadata(self):
        """Test logging with metadata."""
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
        )
        
        request = {"model": "gpt-4", "prompt": "test"}
        response = {"completion": "response"}
        metadata = {"user_id": "123", "session_id": "abc"}
        
        logger.log(request=request, response=response, metadata=metadata)
        
        # Get the queued item
        log_entry = logger.log_queue.get_nowait()
        
        assert log_entry["request"] == request
        assert log_entry["response"] == response
        assert log_entry["metadata"] == metadata
        assert "timestamp" in log_entry
        
        logger.shutdown()
    
    def test_get_stats(self):
        """Test statistics tracking."""
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
        )
        
        stats = logger.get_stats()
        
        assert "logs_queued" in stats
        assert "logs_sent" in stats
        assert "logs_failed" in stats
        assert "batches_sent" in stats
        assert "batches_failed" in stats
        
        logger.shutdown()


class TestThreadSafety:
    """Test thread-safe operations."""
    
    def test_concurrent_logging(self):
        """Test multiple threads can log simultaneously."""
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            max_queue_size=1000,
        )
        
        num_threads = 10
        logs_per_thread = 20
        
        def log_worker(thread_id):
            for i in range(logs_per_thread):
                logger.log(
                    request={"thread": thread_id, "index": i},
                    response={"result": f"thread_{thread_id}_log_{i}"},
                )
        
        threads = []
        for i in range(num_threads):
            thread = threading.Thread(target=log_worker, args=(i,))
            threads.append(thread)
            thread.start()
        
        # Wait for all threads to complete
        for thread in threads:
            thread.join()
        
        expected_logs = num_threads * logs_per_thread
        assert logger.stats["logs_queued"] == expected_logs
        
        logger.shutdown()
    
    def test_queue_overflow_handling(self):
        """Test behavior when queue is full."""
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            max_queue_size=5,  # Very small queue
            batch_size=100,  # Large batch size to prevent flushing
        )
        
        # Fill the queue
        for i in range(5):
            success = logger.log(
                request={"index": i},
                response={"result": i},
            )
            assert success is True
        
        # Queue should be full now
        success = logger.log(
            request={"overflow": True},
            response={"result": "should fail"},
        )
        
        assert success is False
        assert logger.stats["logs_failed"] >= 1
        
        logger.shutdown()


class TestBatching:
    """Test batching behavior."""
    
    @patch('burt_logger.logger.requests.post')
    def test_batch_size_trigger(self, mock_post):
        """Test that batching is triggered by batch size."""
        mock_post.return_value.status_code = 200
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=3,
            flush_interval=100.0,  # Long interval so only size triggers
        )
        
        # Log 3 entries to trigger batch
        for i in range(3):
            logger.log(
                request={"index": i},
                response={"result": i},
            )
        
        # Give worker thread time to process
        time.sleep(0.5)
        
        # Should have sent one batch
        assert mock_post.call_count >= 1
        
        # Check the payload
        call_args = mock_post.call_args
        payload = call_args[1]["json"]
        assert "logs" in payload
        assert len(payload["logs"]) == 3
        
        logger.shutdown()
    
    @patch('burt_logger.logger.requests.post')
    def test_time_interval_trigger(self, mock_post):
        """Test that batching is triggered by time interval."""
        mock_post.return_value.status_code = 200
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=100,  # Large batch size
            flush_interval=1.0,  # Short interval
        )
        
        # Log just 2 entries (less than batch size)
        logger.log(request={"test": 1}, response={"result": 1})
        logger.log(request={"test": 2}, response={"result": 2})
        
        # Wait for flush interval
        time.sleep(1.5)
        
        # Should have sent batch due to time trigger
        assert mock_post.call_count >= 1
        
        logger.shutdown()
    
    @patch('burt_logger.logger.requests.post')
    def test_flush_method(self, mock_post):
        """Test manual flush method."""
        mock_post.return_value.status_code = 200
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=100,
            flush_interval=100.0,
        )
        
        # Log some entries
        for i in range(5):
            logger.log(request={"i": i}, response={"r": i})
        
        # Manually flush
        logger.flush(timeout=2.0)
        
        # Should have sent the batch
        assert mock_post.call_count >= 1
        
        logger.shutdown()


class TestRetryLogic:
    """Test retry and error handling."""
    
    @patch('burt_logger.logger.requests.post')
    def test_successful_send(self, mock_post):
        """Test successful log sending."""
        mock_post.return_value.status_code = 200
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=2,
        )
        
        logger.log(request={"test": 1}, response={"result": 1})
        logger.log(request={"test": 2}, response={"result": 2})
        
        time.sleep(0.5)
        
        assert mock_post.call_count >= 1
        assert logger.stats["batches_sent"] >= 1
        assert logger.stats["logs_sent"] >= 2
        
        logger.shutdown()
    
    @patch('burt_logger.logger.requests.post')
    @patch('burt_logger.logger.time.sleep')  # Speed up test
    def test_retry_on_server_error(self, mock_sleep, mock_post):
        """Test retry logic on 5xx errors."""
        # First two calls fail, third succeeds
        mock_post.side_effect = [
            Mock(status_code=500),
            Mock(status_code=503),
            Mock(status_code=200),
        ]
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=1,
            max_retries=3,
        )
        
        logger.log(request={"test": 1}, response={"result": 1})
        
        time.sleep(0.5)
        
        # Should have retried and eventually succeeded
        assert mock_post.call_count == 3
        assert logger.stats["batches_sent"] >= 1
        
        logger.shutdown()
    
    @patch('burt_logger.logger.requests.post')
    @patch('burt_logger.logger.time.sleep')
    def test_no_retry_on_client_error(self, mock_sleep, mock_post):
        """Test that 4xx errors don't trigger retry."""
        mock_post.return_value = Mock(status_code=400, text="Bad request")
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=1,
            max_retries=3,
        )
        
        logger.log(request={"test": 1}, response={"result": 1})
        
        time.sleep(0.5)
        
        # Should have tried only once (no retries for 4xx)
        assert mock_post.call_count == 1
        assert logger.stats["batches_failed"] >= 1
        
        logger.shutdown()
    
    @patch('burt_logger.logger.requests.post')
    @patch('burt_logger.logger.time.sleep')
    def test_exponential_backoff(self, mock_sleep, mock_post):
        """Test exponential backoff timing."""
        mock_post.return_value = Mock(status_code=500)
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=1,
            max_retries=3,
            initial_retry_delay=1.0,
        )
        
        logger.log(request={"test": 1}, response={"result": 1})
        
        time.sleep(0.5)
        
        # Check that sleep was called with increasing delays
        sleep_calls = [call[0][0] for call in mock_sleep.call_args_list]
        
        # Should have exponential backoff: 1, 2, 4 seconds
        assert len(sleep_calls) >= 2
        if len(sleep_calls) >= 2:
            assert sleep_calls[1] > sleep_calls[0]
        
        logger.shutdown()


class TestShutdown:
    """Test graceful shutdown behavior."""
    
    @patch('burt_logger.logger.requests.post')
    def test_graceful_shutdown_flushes_logs(self, mock_post):
        """Test that shutdown flushes remaining logs."""
        mock_post.return_value.status_code = 200
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=100,  # Large batch to prevent auto-flushing
            flush_interval=100.0,
        )
        
        # Add some logs
        for i in range(5):
            logger.log(request={"i": i}, response={"r": i})
        
        # Shutdown should flush them
        logger.shutdown(timeout=2.0)
        
        # Verify logs were sent
        assert mock_post.call_count >= 1
        
    def test_context_manager(self):
        """Test using logger as context manager."""
        with LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
        ) as logger:
            logger.log(request={"test": 1}, response={"result": 1})
            assert logger._worker_thread.is_alive()
        
        # After exiting context, shutdown should have been called
        # Worker thread should eventually stop
        time.sleep(0.5)
        assert not logger._worker_thread.is_alive() or logger._shutdown.is_set()
    
    def test_multiple_shutdown_calls(self):
        """Test that multiple shutdown calls don't cause errors."""
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
        )
        
        logger.shutdown()
        logger.shutdown()  # Should not raise error
        logger.shutdown()
        
        assert logger._shutdown.is_set()


class TestHTTPHeaders:
    """Test HTTP request formatting."""
    
    @patch('burt_logger.logger.requests.post')
    def test_authorization_header(self, mock_post):
        """Test that API key is sent in Authorization header."""
        mock_post.return_value.status_code = 200
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="sk-test-key-123",
            batch_size=1,
        )
        
        logger.log(request={"test": 1}, response={"result": 1})
        
        time.sleep(0.5)
        
        # Check headers
        call_args = mock_post.call_args
        headers = call_args[1]["headers"]
        
        assert "Authorization" in headers
        assert headers["Authorization"] == "Bearer sk-test-key-123"
        assert headers["Content-Type"] == "application/json"
        
        logger.shutdown()
    
    @patch('burt_logger.logger.requests.post')
    def test_payload_structure(self, mock_post):
        """Test that payload has correct structure."""
        mock_post.return_value.status_code = 200
        
        logger = LLMLogger(
            endpoint="https://test.com/logs",
            api_key="test-key",
            batch_size=2,
        )
        
        logger.log(request={"model": "gpt-4"}, response={"text": "hello"})
        logger.log(request={"model": "gpt-3.5"}, response={"text": "world"})
        
        time.sleep(0.5)
        
        # Check payload structure
        call_args = mock_post.call_args
        payload = call_args[1]["json"]
        
        assert "logs" in payload
        assert "timestamp" in payload
        assert isinstance(payload["logs"], list)
        assert len(payload["logs"]) == 2
        
        # Check log entry structure
        log_entry = payload["logs"][0]
        assert "request" in log_entry
        assert "response" in log_entry
        assert "metadata" in log_entry
        assert "timestamp" in log_entry
        
        logger.shutdown()


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

