"""
Comprehensive test suite for Sector8 SDK to improve coverage
"""

import pytest
import os
import asyncio
from unittest.mock import patch, MagicMock, AsyncMock, Mock
from sector8.simple_client import Sector8SimpleClient
from sector8 import setup, track, get_simple_client, configure_wizard
from sector8.errors import ConfigurationError, NetworkError, AuthenticationError


class TestSector8SDK:
    """Test cases for main SDK module"""

    def test_setup_with_api_key(self):
        """Test setup function with API key"""
        client = setup(api_key="test-key-123456789")
        assert isinstance(client, Sector8SimpleClient)
        assert client.api_key == "test-key-123456789"

    def test_setup_without_api_key(self):
        """Test setup function without API key raises error"""
        with patch.dict(os.environ, {}, clear=True):
            with pytest.raises(ConfigurationError):
                setup()

    def test_get_simple_client_before_setup(self):
        """Test getting client before setup raises error"""
        # Clear global client
        import sector8
        sector8._simple_client = None
        
        with pytest.raises(RuntimeError, match="SDK not initialized"):
            get_simple_client()

    @patch('sector8.simple_client.aiohttp.ClientSession')
    def test_track_decorator_async_function(self, mock_session_class):
        """Test track decorator on async function"""
        
        @track()
        async def async_llm_function(text: str) -> str:
            return f"Processed: {text}"
        
        # Run the decorated function
        result = asyncio.run(async_llm_function("test"))
        assert result == "Processed: test"

    @patch('sector8.simple_client.aiohttp.ClientSession')
    def test_track_decorator_sync_function(self, mock_session_class):
        """Test track decorator on sync function"""
        
        @track()
        def sync_llm_function(text: str) -> str:
            return f"Processed: {text}"
        
        # Run the decorated function
        result = sync_llm_function("test")
        assert result == "Processed: test"

    @patch('sector8.simple_client.aiohttp.ClientSession')
    def test_track_decorator_with_exception(self, mock_session_class):
        """Test track decorator when function raises exception"""
        
        @track()
        def failing_function():
            raise ValueError("Test error")
        
        with pytest.raises(ValueError, match="Test error"):
            failing_function()

    def test_configure_wizard(self):
        """Test configuration wizard import"""
        # Just test that the function can be imported
        from sector8 import configure_wizard
        assert callable(configure_wizard)


class TestSector8SimpleClientComprehensive:
    """Comprehensive tests for Sector8SimpleClient"""

    def test_client_with_custom_config(self):
        """Test client initialization with custom configuration"""
        client = Sector8SimpleClient(
            api_key="test-key-123456789",
            base_url="https://custom-api.sector8.ai",
            client_id="custom-client",
            debug=True
        )
        assert client.api_key == "test-key-123456789"
        assert client.base_url == "https://custom-api.sector8.ai"
        assert client.client_id == "custom-client"
        assert client.debug is True

    @pytest.mark.asyncio
    async def test_save_telemetry_success(self):
        """Test successful telemetry saving"""
        client = Sector8SimpleClient(api_key="test-key-123456789")
        await client.initialize()
        
        # Mock the session's post method directly with proper async context manager
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={'status': 'success'})
        
        # Mock post to return the response as an async context manager
        class MockPost:
            def __init__(self, response):
                self.response = response
            
            async def __aenter__(self):
                return self.response
            
            async def __aexit__(self, exc_type, exc_val, exc_tb):
                pass
        
        def mock_post(*args, **kwargs):
            return MockPost(mock_response)
        
        client.session.post = mock_post
        
        result = await client.save_telemetry(
            provider="openai",
            model="gpt-4",
            prompt="test prompt",
            completion="test completion",
            tokens_used=100,
            latency_ms=1500,
            cost=0.002
        )
        
        assert result == {'status': 'success'}
        await client.close()

    @patch('sector8.simple_client.aiohttp.ClientSession')
    @pytest.mark.asyncio
    async def test_save_threat_alert_success(self, mock_session_class):
        """Test successful threat alert saving"""
        mock_session = AsyncMock()
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={'alert_id': '12345'})
        
        # Create proper async context manager mock
        mock_context_manager = AsyncMock()
        mock_context_manager.__aenter__ = AsyncMock(return_value=mock_response)
        mock_context_manager.__aexit__ = AsyncMock(return_value=None)
        
        mock_session.post = Mock(return_value=mock_context_manager)
        mock_session_class.return_value = mock_session

        client = Sector8SimpleClient(api_key="test-key-123456789")
        result = await client.save_threat_alert(
            threat_type="prompt_injection",
            severity="High",
            status="active",
            description="Test threat"
        )
        
        assert result == {'alert_id': '12345'}

    @patch('sector8.simple_client.aiohttp.ClientSession')
    @pytest.mark.asyncio
    async def test_save_incident_log_success(self, mock_session_class):
        """Test successful incident log saving"""
        mock_session = AsyncMock()
        mock_response = AsyncMock()
        mock_response.status = 200
        mock_response.json = AsyncMock(return_value={'incident_id': '67890'})
        
        # Create proper async context manager mock
        mock_context_manager = AsyncMock()
        mock_context_manager.__aenter__ = AsyncMock(return_value=mock_response)
        mock_context_manager.__aexit__ = AsyncMock(return_value=None)
        
        mock_session.post = Mock(return_value=mock_context_manager)
        mock_session_class.return_value = mock_session

        client = Sector8SimpleClient(api_key="test-key-123456789")
        result = await client.save_incident_log(
            severity="medium",
            tool="guard_analysis",
            classification="pii_detection",
            details="Test details",
            issue="Test issue",
            tokens=100
        )
        
        assert result == {'incident_id': '67890'}

    @patch('sector8.simple_client.aiohttp.ClientSession')
    @pytest.mark.asyncio
    async def test_network_error_handling(self, mock_session_class):
        """Test network error handling with retries"""
        mock_session = AsyncMock()
        mock_session.post = AsyncMock(side_effect=asyncio.TimeoutError("Connection timeout"))
        mock_session_class.return_value = mock_session

        client = Sector8SimpleClient(api_key="test-key-123456789")
        
        with pytest.raises(NetworkError, match="Request timed out"):
            await client.save_telemetry(
                provider="openai",
                model="gpt-4", 
                prompt="test",
                completion="test",
                tokens_used=100,
                latency_ms=1500,
                cost=0.002
            )

    @patch('sector8.simple_client.aiohttp.ClientSession')
    @pytest.mark.asyncio
    async def test_authentication_error_handling(self, mock_session_class):
        """Test authentication error handling"""
        mock_session = AsyncMock()
        mock_response = AsyncMock()
        mock_response.status = 401
        mock_response.text = AsyncMock(return_value="Unauthorized")
        
        # Create proper async context manager mock
        mock_context_manager = AsyncMock()
        mock_context_manager.__aenter__ = AsyncMock(return_value=mock_response)
        mock_context_manager.__aexit__ = AsyncMock(return_value=None)
        
        mock_session.post = Mock(return_value=mock_context_manager)
        mock_session_class.return_value = mock_session

        client = Sector8SimpleClient(api_key="test-key-123456789")
        
        with pytest.raises(AuthenticationError, match="Invalid API key"):
            await client.save_telemetry(
                provider="openai",
                model="gpt-4",
                prompt="test", 
                completion="test",
                tokens_used=100,
                latency_ms=1500,
                cost=0.002
            )

    @patch('sector8.simple_client.aiohttp.ClientSession')
    @pytest.mark.asyncio
    async def test_rate_limiting_retry(self, mock_session_class):
        """Test rate limiting with retry logic"""
        mock_session = AsyncMock()
        
        # First attempt: rate limited, second attempt: success
        mock_response_429 = AsyncMock()
        mock_response_429.status = 429
        
        mock_response_200 = AsyncMock()
        mock_response_200.status = 200
        mock_response_200.json = AsyncMock(return_value={'status': 'success'})
        
        # Create proper async context manager mocks
        mock_context_manager_429 = AsyncMock()
        mock_context_manager_429.__aenter__ = AsyncMock(return_value=mock_response_429)
        mock_context_manager_429.__aexit__ = AsyncMock(return_value=None)
        
        mock_context_manager_200 = AsyncMock()
        mock_context_manager_200.__aenter__ = AsyncMock(return_value=mock_response_200)
        mock_context_manager_200.__aexit__ = AsyncMock(return_value=None)
        
        mock_session.post = AsyncMock(side_effect=[mock_context_manager_429, mock_context_manager_200])
        mock_session_class.return_value = mock_session

        client = Sector8SimpleClient(api_key="test-key-123456789", debug=True)
        
        with patch('asyncio.sleep'):  # Mock sleep to speed up test
            result = await client.save_telemetry(
                provider="openai",
                model="gpt-4",
                prompt="test",
                completion="test", 
                tokens_used=100,
                latency_ms=1500,
                cost=0.002
            )
            
        assert result == {'status': 'success'}

    @pytest.mark.asyncio
    async def test_client_cleanup(self):
        """Test client resource cleanup"""
        client = Sector8SimpleClient(api_key="test-key-123456789")
        await client.initialize()
        
        # Mock session for cleanup test
        mock_session = AsyncMock()
        mock_session.closed = False
        client.session = mock_session
        
        await client.close()
        mock_session.close.assert_called_once()


class TestUtilityFunctions:
    """Test utility functions in the SDK"""

    def test_provider_detection(self):
        """Test LLM provider detection"""
        from sector8 import _detect_provider_from_function
        
        def openai_function():
            from openai import OpenAI
            client = OpenAI()
            return client
            
        def anthropic_function():
            import anthropic
            return anthropic.client
            
        # Test detection (returns None if not found in source)
        result_openai = _detect_provider_from_function(openai_function, "test")
        result_anthropic = _detect_provider_from_function(anthropic_function, "test")
        
        # These may be None depending on function source availability
        assert result_openai in [None, "openai"]
        assert result_anthropic in [None, "anthropic"]

    def test_token_estimation(self):
        """Test token estimation function"""
        from sector8 import _estimate_tokens_from_result
        
        result = _estimate_tokens_from_result("This is a test string")
        assert isinstance(result, int)
        assert result >= 1
        
        result_empty = _estimate_tokens_from_result("")
        assert result_empty >= 1

    def test_cost_estimation(self):
        """Test cost estimation function"""
        from sector8 import _estimate_cost_from_result
        
        cost = _estimate_cost_from_result("test response", "openai", "gpt-4")
        assert isinstance(cost, float)
        assert cost >= 0
        
        cost_unknown = _estimate_cost_from_result("test", "unknown", "unknown")
        assert cost_unknown >= 0

    def test_prompt_extraction(self):
        """Test prompt extraction from arguments"""
        from sector8 import _extract_prompt_from_args
        
        # Test with keyword arguments
        prompt = _extract_prompt_from_args([], {'text': 'test prompt'})
        assert prompt == 'test prompt'
        
        # Test with positional arguments
        prompt = _extract_prompt_from_args(['long test prompt'], {})
        assert prompt == 'long test prompt'
        
        # Test with no valid prompt
        prompt = _extract_prompt_from_args(['short'], {})
        assert prompt is None

    def test_completion_extraction(self):
        """Test completion extraction from results"""
        from sector8 import _extract_completion_from_result
        
        completion = _extract_completion_from_result("test completion response")
        assert completion == "test completion response"
        
        completion_none = _extract_completion_from_result(123)
        assert completion_none is None


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