"""
Unit tests for response_generator module
Tests LLM response generation functionality
"""

import unittest
from unittest.mock import MagicMock, Mock, patch

from src.response_generator import (_generate_response_impl,
                                    format_chat_history,
                                    generate_xeno_response)


class TestResponseGenerator(unittest.TestCase):
    """Test cases for response_generator module"""

    def setUp(self):
        """Set up test fixtures"""
        self.context = """Knowledge Entry 1:
Q: How do I create an account?
A: Visit our website and click Sign Up.
----------------------------------------"""

        self.question = "How can I create an account?"

        self.chat_history = [
            {"role": "user", "content": "Hello"},
            {"role": "assistant", "content": "Hi! How can I help you?"},
        ]

    def test_format_chat_history(self):
        """Test formatting chat history"""
        formatted = format_chat_history(self.chat_history)

        # Check format
        self.assertIn("User: Hello", formatted)
        self.assertIn("Assistant: Hi! How can I help you?", formatted)
        self.assertIn("\n", formatted)

    def test_format_chat_history_empty(self):
        """Test formatting empty chat history"""
        formatted = format_chat_history([])
        self.assertEqual(formatted, "No previous conversation")

    def test_format_chat_history_single_message(self):
        """Test formatting single message"""
        history = [{"role": "user", "content": "Hello"}]
        formatted = format_chat_history(history)
        self.assertEqual(formatted, "User: Hello")

    def test_format_chat_history_missing_fields(self):
        """Test formatting with missing fields"""
        history = [
            {"role": "user"},  # Missing content
            {"content": "Test"},  # Missing role
        ]
        formatted = format_chat_history(history)
        self.assertIn("User:", formatted)
        self.assertIn("Unknown:", formatted)

    @patch("src.response_generator.genai_client")
    def test_generate_response_impl(self, mock_genai_client):
        """Test internal response generation implementation"""
        # Configure mock response
        mock_genai_client.models.generate_content.return_value.text = "You can create an account by visiting our website."

        response = _generate_response_impl(
            self.context, self.question, self.chat_history
        )

        # Verify generate_content was called with model and content
        mock_genai_client.models.generate_content.assert_called_once()
        call_kwargs = mock_genai_client.models.generate_content.call_args[1]
        self.assertIn("model", call_kwargs)
        self.assertIn("contents", call_kwargs)

        # Check response
        self.assertEqual(response, "You can create an account by visiting our website.")

    @patch("src.response_generator.genai_client")
    def test_generate_response_with_empty_history(self, mock_genai_client):
        """Test generating response with empty history"""
        mock_genai_client.models.generate_content.return_value.text = "Test response"

        response = _generate_response_impl(self.context, self.question, [])

        # Verify it still works
        self.assertEqual(response, "Test response")

        # Check that "None" was used for history in prompt
        call_kwargs = mock_genai_client.models.generate_content.call_args[1]
        prompt = call_kwargs["contents"]
        self.assertIn("None", prompt)

    @patch("src.response_generator.genai_client")
    def test_prompt_structure(self, mock_genai_client):
        """Test that prompt includes all necessary components"""
        mock_genai_client.models.generate_content.return_value.text = "Test response"

        _generate_response_impl(self.context, self.question, self.chat_history)

        # Get the prompt that was sent
        call_kwargs = mock_genai_client.models.generate_content.call_args[1]
        prompt = call_kwargs["contents"]

        # Verify prompt structure
        self.assertIn("HISTORY", prompt)
        self.assertIn("CONTEXT", prompt)
        self.assertIn("QUESTION", prompt)
        self.assertIn(self.context, prompt)
        self.assertIn(self.question, prompt)

    @patch("src.response_generator.genai_client")
    def test_generate_xeno_response_with_timer(self, mock_genai_client):
        """Test generate_xeno_response with timer"""
        mock_genai_client.models.generate_content.return_value.text = "Test response"

        mock_timer = Mock()
        mock_timer.time_step = MagicMock()
        mock_timer.time_step.return_value.__enter__ = Mock()
        mock_timer.time_step.return_value.__exit__ = Mock()

        response = generate_xeno_response(
            self.context, self.question, self.chat_history, timer=mock_timer
        )

        # Verify timer was used
        mock_timer.time_step.assert_called_once_with("llm_generation")

        # Verify response
        self.assertEqual(response, "Test response")

    @patch("src.response_generator.genai_client")
    def test_response_text_stripping(self, mock_genai_client):
        """Test that response text is stripped of whitespace"""
        mock_genai_client.models.generate_content.return_value.text = "Test response with spaces"

        response = _generate_response_impl(self.context, self.question, [])

        # Response should be returned as-is from mock
        self.assertEqual(response, "Test response with spaces")

    @patch("src.response_generator.genai_client")
    def test_system_prompt_inclusion(self, mock_genai_client):
        """Test that system prompt is included in generated prompt"""
        mock_genai_client.models.generate_content.return_value.text = "Test"

        _generate_response_impl(self.context, self.question, [])

        # Get the prompt
        call_kwargs = mock_genai_client.models.generate_content.call_args[1]
        prompt = call_kwargs["contents"]

        # Should contain system prompt text
        self.assertIn("XENO Support Assistant", prompt)


if __name__ == "__main__":
    unittest.main()