"""
Tests for q_business_chat_cli.py

This test suite covers:
- Message sending functionality
- Session management
- Error handling
- Response parsing
"""

import pytest
from unittest.mock import Mock, patch, MagicMock
import json
import sys
from io import StringIO


# Mock the boto3 client before importing the module
@pytest.fixture(autouse=True)
def mock_boto3():
    """Mock boto3 client for all tests"""
    with patch('boto3.client') as mock_client:
        yield mock_client


class TestSendMessageToQ:
    """Test suite for send_message_to_q function"""

    def test_send_message_without_session_id(self, mock_boto3):
        """Test sending a message without an existing session ID"""
        # Import after mocking
        import q_business_chat_cli

        # Setup mock response
        mock_qbusiness = Mock()
        mock_response = {
            'conversationId': 'conv-123',
            'messages': [
                {
                    'type': 'ASSISTANT_MESSAGE',
                    'payload': {'text': 'Hello! How can I help you?'}
                }
            ]
        }
        mock_qbusiness.chat_sync.return_value = mock_response
        mock_boto3.return_value = mock_qbusiness

        # Reinitialize the client in the module
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        # Test the function
        response = q_business_chat_cli.send_message_to_q("Hello", None)

        # Assertions
        assert response == mock_response
        assert response['conversationId'] == 'conv-123'
        mock_qbusiness.chat_sync.assert_called_once()
        call_args = mock_qbusiness.chat_sync.call_args[1]
        assert 'conversationId' not in call_args
        assert call_args['applicationId'] == q_business_chat_cli.APPLICATION_ID
        assert call_args['userId'] == q_business_chat_cli.USER_ID

    def test_send_message_with_session_id(self, mock_boto3):
        """Test sending a message with an existing session ID"""
        import q_business_chat_cli

        # Setup mock
        mock_qbusiness = Mock()
        mock_response = {
            'conversationId': 'conv-123',
            'messages': [
                {
                    'type': 'ASSISTANT_MESSAGE',
                    'payload': {'text': 'Follow-up response'}
                }
            ]
        }
        mock_qbusiness.chat_sync.return_value = mock_response
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        # Test with session ID
        response = q_business_chat_cli.send_message_to_q("Follow-up question", "conv-123")

        # Assertions
        assert response == mock_response
        call_args = mock_qbusiness.chat_sync.call_args[1]
        assert call_args['conversationId'] == 'conv-123'

    def test_send_message_with_error(self, mock_boto3):
        """Test error handling when API call fails"""
        import q_business_chat_cli

        # Setup mock to raise exception
        mock_qbusiness = Mock()
        mock_qbusiness.chat_sync.side_effect = Exception("API Error")
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        # Capture print output
        captured_output = StringIO()

        with patch('sys.stdout', captured_output):
            response = q_business_chat_cli.send_message_to_q("Test", None)

        # Assertions
        assert response is None
        assert "Error communicating with Amazon Q Business" in captured_output.getvalue()

    def test_message_payload_format(self, mock_boto3):
        """Test that message payload is correctly formatted"""
        import q_business_chat_cli

        # Setup mock
        mock_qbusiness = Mock()
        mock_qbusiness.chat_sync.return_value = {'conversationId': 'conv-123', 'messages': []}
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        # Send message
        test_message = "What is the weather?"
        q_business_chat_cli.send_message_to_q(test_message, None)

        # Verify payload structure
        call_args = mock_qbusiness.chat_sync.call_args[1]
        messages = call_args['messages']

        assert len(messages) == 1
        assert messages[0]['type'] == 'USER_MESSAGE'
        assert messages[0]['source'] == 'CLI_CHAT'
        assert messages[0]['payload']['text'] == test_message


class TestResponseParsing:
    """Test suite for response parsing logic"""

    def test_extract_assistant_message(self):
        """Test extracting assistant message from response"""
        response = {
            'conversationId': 'conv-123',
            'messages': [
                {
                    'type': 'USER_MESSAGE',
                    'payload': {'text': 'Question'}
                },
                {
                    'type': 'ASSISTANT_MESSAGE',
                    'payload': {'text': 'Answer'}
                }
            ]
        }

        # Simulate the logic from the main loop
        q_message = next((msg for msg in response.get('messages', [])
                         if msg.get('type') == 'ASSISTANT_MESSAGE'), None)

        assert q_message is not None
        assert q_message['payload']['text'] == 'Answer'

    def test_extract_message_when_missing(self):
        """Test handling when assistant message is missing"""
        response = {
            'conversationId': 'conv-123',
            'messages': []
        }

        q_message = next((msg for msg in response.get('messages', [])
                         if msg.get('type') == 'ASSISTANT_MESSAGE'), None)

        assert q_message is None

    def test_source_attributions_parsing(self):
        """Test parsing source attributions"""
        response = {
            'conversationId': 'conv-123',
            'messages': [
                {
                    'type': 'ASSISTANT_MESSAGE',
                    'payload': {'text': 'Based on the documentation...'}
                }
            ],
            'sourceAttributions': [
                {
                    'title': 'User Guide',
                    'url': 'https://example.com/guide'
                },
                {
                    'title': 'FAQ',
                    'url': 'https://example.com/faq'
                }
            ]
        }

        assert 'sourceAttributions' in response
        assert len(response['sourceAttributions']) == 2
        assert response['sourceAttributions'][0]['title'] == 'User Guide'
        assert response['sourceAttributions'][1]['url'] == 'https://example.com/faq'


class TestConfiguration:
    """Test suite for configuration validation"""

    def test_default_configuration(self):
        """Test that default configuration values are set"""
        import q_business_chat_cli

        assert q_business_chat_cli.APPLICATION_ID is not None
        assert q_business_chat_cli.REGION is not None
        assert q_business_chat_cli.USER_ID is not None

    def test_region_configuration(self):
        """Test region configuration"""
        import q_business_chat_cli

        # Verify region is a valid AWS region format
        region = q_business_chat_cli.REGION
        assert isinstance(region, str)
        assert len(region) > 0


class TestSessionManagement:
    """Test suite for session ID management"""

    def test_session_id_persistence(self, mock_boto3):
        """Test that session ID persists across messages"""
        import q_business_chat_cli

        # Setup mock
        mock_qbusiness = Mock()
        responses = [
            {'conversationId': 'conv-abc', 'messages': []},
            {'conversationId': 'conv-abc', 'messages': []},
        ]
        mock_qbusiness.chat_sync.side_effect = responses
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        # First message
        response1 = q_business_chat_cli.send_message_to_q("First message", None)
        session_id = response1.get('conversationId')

        # Second message with session ID
        response2 = q_business_chat_cli.send_message_to_q("Second message", session_id)

        # Verify second call included the session ID
        second_call_args = mock_qbusiness.chat_sync.call_args[1]
        assert second_call_args['conversationId'] == 'conv-abc'


class TestErrorScenarios:
    """Test suite for various error scenarios"""

    def test_empty_response(self, mock_boto3):
        """Test handling of empty response"""
        import q_business_chat_cli

        mock_qbusiness = Mock()
        mock_qbusiness.chat_sync.return_value = {}
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        response = q_business_chat_cli.send_message_to_q("Test", None)

        assert response == {}
        assert 'conversationId' not in response

    def test_malformed_response(self, mock_boto3):
        """Test handling of malformed response"""
        import q_business_chat_cli

        mock_qbusiness = Mock()
        mock_qbusiness.chat_sync.return_value = {
            'messages': [
                {
                    'type': 'ASSISTANT_MESSAGE',
                    # Missing 'payload' field
                }
            ]
        }
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        response = q_business_chat_cli.send_message_to_q("Test", None)

        # Should return the response even if malformed
        assert response is not None
        assert 'messages' in response

    def test_network_timeout(self, mock_boto3):
        """Test handling of network timeout"""
        import q_business_chat_cli

        mock_qbusiness = Mock()
        mock_qbusiness.chat_sync.side_effect = Exception("Request timeout")
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        response = q_business_chat_cli.send_message_to_q("Test", None)

        assert response is None


class TestIntegration:
    """Integration tests simulating full conversation flow"""

    def test_multi_turn_conversation(self, mock_boto3):
        """Test a multi-turn conversation flow"""
        import q_business_chat_cli

        mock_qbusiness = Mock()

        # Simulate multiple exchanges
        mock_responses = [
            {
                'conversationId': 'conv-xyz',
                'messages': [{'type': 'ASSISTANT_MESSAGE', 'payload': {'text': 'Response 1'}}]
            },
            {
                'conversationId': 'conv-xyz',
                'messages': [{'type': 'ASSISTANT_MESSAGE', 'payload': {'text': 'Response 2'}}]
            },
            {
                'conversationId': 'conv-xyz',
                'messages': [{'type': 'ASSISTANT_MESSAGE', 'payload': {'text': 'Response 3'}}]
            }
        ]
        mock_qbusiness.chat_sync.side_effect = mock_responses
        q_business_chat_cli.qbusiness_client = mock_qbusiness

        # Simulate conversation
        session_id = None

        # Turn 1
        response1 = q_business_chat_cli.send_message_to_q("Question 1", session_id)
        session_id = response1.get('conversationId')
        assert session_id == 'conv-xyz'

        # Turn 2
        response2 = q_business_chat_cli.send_message_to_q("Question 2", session_id)
        assert response2.get('conversationId') == 'conv-xyz'

        # Turn 3
        response3 = q_business_chat_cli.send_message_to_q("Question 3", session_id)
        assert response3.get('conversationId') == 'conv-xyz'

        # Verify all calls were made
        assert mock_qbusiness.chat_sync.call_count == 3


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