"""
Tests for EntityAwareQuestionGenerator.

Tests cover:
- Entity detection with various keywords
- Entity-focused prompt generation with/without past questions
- Response parsing for Q: prefixed lines
- Non-entity query fallback to parent class
- Sub-question generation for entity queries
- Error handling for exceptions
"""

from datetime import datetime, UTC
from unittest.mock import Mock, patch

import pytest

from local_deep_research.advanced_search_system.questions.entity_aware_question import (
    EntityAwareQuestionGenerator,
)


class TestEntityAwareQuestionGeneratorInit:
    """Tests for EntityAwareQuestionGenerator initialization."""

    def test_init_stores_model(self):
        """Generator stores the model reference."""
        mock_model = Mock()
        generator = EntityAwareQuestionGenerator(mock_model)
        assert generator.model is mock_model


class TestEntityDetection:
    """Tests for entity keyword detection."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(content="Q: query1\nQ: query2")
        return mock

    @pytest.fixture
    def generator(self, mock_model):
        """Create generator instance."""
        return EntityAwareQuestionGenerator(mock_model)

    def test_detect_who_keyword(self, generator):
        """Detect 'who' as entity keyword."""
        generator.generate_questions("", "Who is the president?", 2)
        # Entity query uses entity-focused prompt, check model was called
        assert generator.model.invoke.called
        call_args = generator.model.invoke.call_args[0][0]
        assert "entity" in call_args.lower() or "identify" in call_args.lower()

    def test_detect_what_keyword(self, generator):
        """Detect 'what' as entity keyword."""
        generator.generate_questions("", "What character is this?", 2)
        assert generator.model.invoke.called

    def test_detect_which_keyword(self, generator):
        """Detect 'which' as entity keyword."""
        generator.generate_questions("", "Which movie won the award?", 2)
        assert generator.model.invoke.called

    def test_detect_identify_keyword(self, generator):
        """Detect 'identify' as entity keyword."""
        generator.generate_questions("", "Identify the person", 2)
        assert generator.model.invoke.called

    def test_detect_name_keyword(self, generator):
        """Detect 'name' as entity keyword."""
        generator.generate_questions("", "Name the scientist", 2)
        assert generator.model.invoke.called

    def test_detect_character_keyword(self, generator):
        """Detect 'character' as entity keyword."""
        generator.generate_questions("", "The character from TV", 2)
        assert generator.model.invoke.called

    def test_detect_person_keyword(self, generator):
        """Detect 'person' as entity keyword."""
        generator.generate_questions("", "The person responsible", 2)
        assert generator.model.invoke.called

    def test_detect_company_keyword(self, generator):
        """Detect 'company' as entity keyword."""
        generator.generate_questions("", "The company that made it", 2)
        assert generator.model.invoke.called

    def test_detect_author_keyword(self, generator):
        """Detect 'author' as entity keyword."""
        generator.generate_questions("", "The author of the book", 2)
        assert generator.model.invoke.called

    def test_case_insensitive_detection(self, generator):
        """Entity detection is case-insensitive."""
        generator.generate_questions("", "WHO is the president?", 2)
        assert generator.model.invoke.called


class TestNonEntityQueryFallback:
    """Tests for non-entity query fallback to parent class."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(content="Q: query1\nQ: query2")
        return mock

    @pytest.fixture
    def generator(self, mock_model):
        """Create generator instance."""
        return EntityAwareQuestionGenerator(mock_model)

    def test_non_entity_query_returns_empty(self, generator):
        """Non-entity query returns empty list without calling the model."""
        result = generator.generate_questions(
            "", "How to learn programming?", 2
        )
        assert result == []
        generator.model.invoke.assert_not_called()

    def test_technical_query_uses_parent(self, generator):
        """Technical queries use parent class."""
        with patch.object(
            EntityAwareQuestionGenerator.__bases__[0],
            "generate_questions",
            return_value=["q1", "q2"],
        ):
            generator.generate_questions(
                "", "Explain machine learning algorithms", 2
            )


class TestEntityFocusedPromptGeneration:
    """Tests for entity-focused prompt generation."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model."""
        mock = Mock()
        mock.invoke.return_value = Mock(content="Q: query1\nQ: query2")
        return mock

    @pytest.fixture
    def generator(self, mock_model):
        """Create generator instance."""
        return EntityAwareQuestionGenerator(mock_model)

    def test_prompt_includes_query(self, generator):
        """Prompt includes the original query."""
        generator.generate_questions("", "Who is the main character?", 2)
        call_args = generator.model.invoke.call_args[0][0]
        assert "Who is the main character?" in call_args

    def test_prompt_includes_current_date(self, generator):
        """Prompt includes current date."""
        generator.generate_questions("", "Who is the president?", 2)
        call_args = generator.model.invoke.call_args[0][0]
        # Should contain date in YYYY-MM-DD format
        today = datetime.now(UTC).strftime("%Y-%m-%d")
        assert today in call_args

    def test_prompt_with_past_questions(self, generator):
        """Prompt includes past questions when provided."""
        past = {1: ["Q1", "Q2"], 2: ["Q3"]}
        generator.generate_questions("", "Who is the character?", 2, past)
        call_args = generator.model.invoke.call_args[0][0]
        assert "Past questions:" in call_args

    def test_prompt_without_past_questions(self, generator):
        """Prompt format changes when no past questions."""
        generator.generate_questions("", "Who is the inventor?", 2, None)
        call_args = generator.model.invoke.call_args[0][0]
        # Without past questions, uses simpler format
        assert "Past questions:" not in call_args

    def test_prompt_includes_current_knowledge(self, generator):
        """Prompt includes current knowledge when with past questions."""
        past = {1: ["Q1"]}
        knowledge = "Some gathered knowledge"
        generator.generate_questions(
            knowledge, "Who is the scientist?", 2, past
        )
        call_args = generator.model.invoke.call_args[0][0]
        assert "Some gathered knowledge" in call_args

    def test_questions_per_iteration_in_prompt(self, generator):
        """Prompt requests specific number of questions."""
        generator.generate_questions("", "Who invented it?", 3)
        call_args = generator.model.invoke.call_args[0][0]
        assert "3" in call_args


class TestResponseParsing:
    """Tests for parsing LLM responses."""

    @pytest.fixture
    def generator(self):
        """Create generator with mock model."""
        mock_model = Mock()
        return EntityAwareQuestionGenerator(mock_model)

    def test_parse_q_prefixed_lines(self, generator):
        """Parse lines starting with 'Q:'."""
        generator.model.invoke.return_value = Mock(
            content="Q: First query here\nQ: Second query here\nQ: Third query"
        )
        result = generator.generate_questions("", "Who is it?", 3)
        assert len(result) == 3
        assert "First query here" in result[0]
        assert "Second query here" in result[1]
        assert "Third query" in result[2]

    def test_respect_questions_per_iteration(self, generator):
        """Limit results to questions_per_iteration."""
        generator.model.invoke.return_value = Mock(
            content="Q: q1\nQ: q2\nQ: q3\nQ: q4\nQ: q5"
        )
        result = generator.generate_questions("", "Who is it?", 2)
        assert len(result) == 2

    def test_skip_non_q_lines(self, generator):
        """Skip lines not starting with 'Q:'."""
        generator.model.invoke.return_value = Mock(
            content="Introduction text\nQ: First query\nMore text\nQ: Second query"
        )
        result = generator.generate_questions("", "Who is it?", 5)
        assert len(result) == 2

    def test_handle_content_attribute(self, generator):
        """Handle response with .content attribute."""
        mock_response = Mock()
        mock_response.content = "Q: query1\nQ: query2"
        generator.model.invoke.return_value = mock_response
        result = generator.generate_questions("", "Who is it?", 2)
        assert len(result) == 2

    def test_handle_string_response(self, generator):
        """Handle plain string response."""
        generator.model.invoke.return_value = "Q: query1\nQ: query2"
        result = generator.generate_questions("", "Who is it?", 2)
        assert len(result) == 2

    def test_strip_q_prefix(self, generator):
        """Strip 'Q:' prefix from results."""
        generator.model.invoke.return_value = Mock(
            content="Q: simple query here"
        )
        result = generator.generate_questions("", "Who is it?", 1)
        assert result[0] == "simple query here"

    def test_strip_whitespace(self, generator):
        """Strip whitespace from parsed queries."""
        generator.model.invoke.return_value = Mock(
            content="Q:   query with spaces   \nQ:  another query  "
        )
        result = generator.generate_questions("", "Who is it?", 2)
        assert result[0] == "query with spaces"
        assert result[1] == "another query"


class TestGenerateSubQuestions:
    """Tests for generate_sub_questions method."""

    @pytest.fixture
    def generator(self):
        """Create generator with mock model."""
        mock_model = Mock()
        return EntityAwareQuestionGenerator(mock_model)

    def test_entity_query_sub_questions(self, generator):
        """Generate sub-questions for entity queries."""
        generator.model.invoke.return_value = Mock(
            content="1. First sub-question\n2. Second sub-question\n3. Third sub-question"
        )
        result = generator.generate_sub_questions("Who invented the telephone?")
        assert len(result) == 3

    def test_parse_numbered_format(self, generator):
        """Parse numbered list format."""
        generator.model.invoke.return_value = Mock(
            content="1. Question one here\n2. Question two here"
        )
        result = generator.generate_sub_questions("Which movie is this?")
        assert "Question one here" in result[0]
        assert "Question two here" in result[1]

    def test_parse_bulleted_format(self, generator):
        """Parse bulleted list format."""
        generator.model.invoke.return_value = Mock(
            content="- First question\n- Second question"
        )
        result = generator.generate_sub_questions("Who wrote this book?")
        assert len(result) == 2

    def test_context_included_in_prompt(self, generator):
        """Context is included in prompt."""
        generator.model.invoke.return_value = Mock(content="1. Q1\n2. Q2")
        generator.generate_sub_questions(
            "Who is the character?",
            context="Previous findings about the character",
        )
        call_args = generator.model.invoke.call_args[0][0]
        assert "Previous findings about the character" in call_args

    def test_non_entity_raises_or_returns_empty(self, generator):
        """Non-entity query attempts parent call (may fail without implementation)."""
        # The base class doesn't have generate_sub_questions, so this tests
        # that the code path is attempted for non-entity queries
        try:
            result = generator.generate_sub_questions(
                "How does machine learning work?"
            )
            # If it doesn't raise, should return empty or valid list
            assert isinstance(result, list)
        except AttributeError:
            # Expected if parent class doesn't implement generate_sub_questions
            pass

    def test_handle_content_attribute(self, generator):
        """Handle response with .content attribute."""
        mock_response = Mock()
        mock_response.content = "1. Sub Q1\n2. Sub Q2"
        generator.model.invoke.return_value = mock_response
        result = generator.generate_sub_questions("Who is the scientist?")
        assert len(result) == 2

    def test_handle_string_response(self, generator):
        """Handle plain string response."""
        generator.model.invoke.return_value = "1. Sub Q1\n2. Sub Q2"
        result = generator.generate_sub_questions("Who is the author?")
        assert len(result) == 2


class TestErrorHandling:
    """Tests for error handling."""

    @pytest.fixture
    def generator(self):
        """Create generator with mock model."""
        mock_model = Mock()
        return EntityAwareQuestionGenerator(mock_model)

    def test_sub_questions_exception_returns_empty(self, generator):
        """Exception in generate_sub_questions returns empty list."""
        generator.model.invoke.side_effect = RuntimeError("Connection failed")
        result = generator.generate_sub_questions("Who is the inventor?")
        assert result == []

    def test_sub_questions_logs_exception(self, generator, caplog):
        """Exception is logged in generate_sub_questions."""
        generator.model.invoke.side_effect = ValueError("Invalid input")
        result = generator.generate_sub_questions("Who is the character?")
        # Exception should be logged
        assert result == []


class TestEdgeCases:
    """Tests for edge cases."""

    @pytest.fixture
    def generator(self):
        """Create generator with mock model."""
        mock_model = Mock()
        mock_model.invoke.return_value = Mock(content="Q: query")
        return EntityAwareQuestionGenerator(mock_model)

    def test_empty_response(self, generator):
        """Handle empty LLM response."""
        generator.model.invoke.return_value = Mock(content="")
        result = generator.generate_questions("", "Who is it?", 2)
        assert result == []

    def test_no_valid_q_lines(self, generator):
        """Handle response with no valid Q: lines."""
        generator.model.invoke.return_value = Mock(
            content="Some text\nMore text\nNo Q lines"
        )
        result = generator.generate_questions("", "Who is it?", 2)
        assert result == []

    def test_empty_query(self, generator):
        """Handle empty query string."""
        generator.generate_questions("", "", 2)
        # Empty query has no entity keywords, should fall back to parent
        # Just verify it doesn't crash

    def test_questions_by_iteration_none(self, generator):
        """Handle questions_by_iteration as None."""
        result = generator.generate_questions("", "Who is it?", 2, None)
        assert isinstance(result, list)

    def test_questions_by_iteration_empty(self, generator):
        """Handle questions_by_iteration as empty dict."""
        result = generator.generate_questions("", "Who is it?", 2, {})
        assert isinstance(result, list)
        # Empty dict is falsy, so uses simpler prompt format

    def test_sub_questions_empty_content(self, generator):
        """Handle empty content in sub-questions."""
        generator.model.invoke.return_value = Mock(content="")
        result = generator.generate_sub_questions("Who is the inventor?")
        assert result == []

    def test_sub_questions_no_numbered_lines(self, generator):
        """Handle response with no numbered lines."""
        generator.model.invoke.return_value = Mock(
            content="Some text without numbers"
        )
        result = generator.generate_sub_questions("Who is the character?")
        assert result == []


class TestIntegration:
    """Integration tests."""

    @pytest.fixture
    def mock_model(self):
        """Create mock model with realistic responses."""
        mock = Mock()
        mock.invoke.return_value = Mock(
            content="""Q: "fictional character" "fourth wall" TV show
Q: character name sitcom breaks fourth wall
Q: TV character speaks to audience directly"""
        )
        return mock

    @pytest.fixture
    def generator(self, mock_model):
        """Create generator instance."""
        return EntityAwareQuestionGenerator(mock_model)

    def test_full_workflow_entity_query(self, generator):
        """Test complete workflow with entity query."""
        result = generator.generate_questions(
            "Some knowledge about characters",
            "What fictional character breaks the fourth wall?",
            3,
            {1: ["Previous question"]},
        )
        assert isinstance(result, list)
        assert len(result) == 3
        assert all(isinstance(q, str) for q in result)

    def test_full_workflow_sub_questions(self, generator):
        """Test complete workflow for sub-questions."""
        generator.model.invoke.return_value = Mock(
            content="1. Who created this character?\n2. What show features this character?\n3. When did the show air?"
        )
        result = generator.generate_sub_questions(
            "Which character breaks the fourth wall?",
            context="Looking for a TV character",
        )
        assert isinstance(result, list)
        assert len(result) == 3
