"""
Tests for advanced_search_system/questions/flexible_browsecomp_question.py

Tests cover:
- FlexibleBrowseCompQuestionGenerator initialization
- _generate_progressive_searches method
- Fallback behavior when not enough searches generated
"""

from unittest.mock import Mock, patch


class TestFlexibleBrowseCompQuestionGeneratorInit:
    """Tests for FlexibleBrowseCompQuestionGenerator initialization."""

    def test_inherits_from_browsecomp(self):
        """Test that FlexibleBrowseCompQuestionGenerator inherits from BrowseCompQuestionGenerator."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        assert issubclass(
            FlexibleBrowseCompQuestionGenerator, BrowseCompQuestionGenerator
        )

    def test_init_with_model(self):
        """Test initialization with a model."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        generator = FlexibleBrowseCompQuestionGenerator(mock_model)

        assert generator.model is mock_model


class TestGenerateProgressiveSearches:
    """Tests for _generate_progressive_searches method."""

    def test_generates_searches_from_model_response(self):
        """Test that searches are generated from model response."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="search query 1\nsearch query 2\nsearch query 3"
        )

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {
            "names": ["John"],
            "temporal": ["2024"],
            "descriptors": ["famous"],
        }

        searches = generator._generate_progressive_searches(
            query="Test query",
            current_knowledge="Some knowledge",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=3,
            iteration=1,
        )

        assert len(searches) == 3

    def test_strips_common_prefixes(self):
        """Test that common prefixes are stripped from searches."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="Q: first query\n- second query\n1. third query"
        )

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {"names": [], "temporal": [], "descriptors": []}

        searches = generator._generate_progressive_searches(
            query="Test",
            current_knowledge="",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=3,
            iteration=1,
        )

        # Should have stripped prefixes
        assert "Q:" not in searches[0] if searches else True
        assert not searches[0].startswith("-") if searches else True

    def test_filters_short_lines(self):
        """Test that short lines are filtered out."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="ab\nshort\nthis is a valid search query\nanother valid query\nthird valid query"
        )

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {
            "names": [],
            "temporal": [],
            "descriptors": [],
            "locations": [],
        }

        searches = generator._generate_progressive_searches(
            query="Test",
            current_knowledge="",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=3,
            iteration=1,
        )

        # "ab" should be filtered (len <= 5)
        assert "ab" not in searches

    def test_filters_lines_ending_with_colon(self):
        """Test that lines ending with colon are filtered."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="Search queries:\nactual search query here\nanother valid search"
        )

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {
            "names": [],
            "temporal": [],
            "descriptors": [],
            "locations": [],
        }

        searches = generator._generate_progressive_searches(
            query="Test",
            current_knowledge="",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=2,
            iteration=1,
        )

        # "Search queries:" should be filtered
        assert not any(s.endswith(":") for s in searches)

    def test_detects_failing_searches(self):
        """Test that failing searches are detected."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(content="broader search query")

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {"names": [], "temporal": [], "descriptors": []}

        # Simulate 5 iterations with 3+ returning 0 results
        results_by_iteration = {1: 0, 2: 0, 3: 0, 4: 5, 5: 0}

        generator._generate_progressive_searches(
            query="Test",
            current_knowledge="",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration=results_by_iteration,
            num_questions=1,
            iteration=6,
        )

        # Check that the prompt mentions failing searches
        call_args = mock_model.invoke.call_args[0][0]
        assert "0 results" in call_args or "broader" in call_args.lower()

    def test_includes_entities_in_prompt(self):
        """Test that entities are included in the prompt."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(content="search query")

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {
            "names": ["Einstein", "Newton"],
            "temporal": ["1905"],
            "descriptors": ["physicist"],
        }

        generator._generate_progressive_searches(
            query="Test",
            current_knowledge="",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=1,
            iteration=1,
        )

        call_args = mock_model.invoke.call_args[0][0]
        assert "Einstein" in call_args
        assert "1905" in call_args

    def test_respects_knowledge_truncate_length(self):
        """Test that knowledge is truncated if truncate length is set."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(content="search query here")

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        generator.knowledge_truncate_length = 100

        entities = {
            "names": [],
            "temporal": [],
            "descriptors": [],
            "locations": [],
        }
        long_knowledge = "A" * 500

        generator._generate_progressive_searches(
            query="Test",
            current_knowledge=long_knowledge,
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=1,
            iteration=1,
        )

        call_args = mock_model.invoke.call_args[0][0]
        # Should be truncated - knowledge_truncate_length limits the knowledge portion
        assert call_args.count("A") < 500  # Definitely truncated from original

    def test_limits_output_to_num_questions(self):
        """Test that output is limited to num_questions."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="query 1\nquery 2\nquery 3\nquery 4\nquery 5"
        )

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {"names": [], "temporal": [], "descriptors": []}

        searches = generator._generate_progressive_searches(
            query="Test",
            current_knowledge="",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=2,
            iteration=1,
        )

        assert len(searches) <= 2

    def test_handles_response_without_content_attr(self):
        """Test handling response without content attribute."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        # Return something that doesn't have content attribute
        mock_response = "search query from string"
        mock_model.invoke.return_value = mock_response

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {"names": [], "temporal": [], "descriptors": []}

        searches = generator._generate_progressive_searches(
            query="Test",
            current_knowledge="",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=1,
            iteration=1,
        )

        # Should handle string response
        assert isinstance(searches, list)

    def test_strips_bullet_prefixes(self):
        """Test that various bullet prefixes are stripped."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="* bullet query\n• unicode bullet\nSearch: prefixed search"
        )

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {"names": [], "temporal": [], "descriptors": []}

        searches = generator._generate_progressive_searches(
            query="Test",
            current_knowledge="",
            entities=entities,
            questions_by_iteration={},
            results_by_iteration={},
            num_questions=3,
            iteration=1,
        )

        # Bullets should be stripped
        for search in searches:
            assert not search.startswith("*")
            assert not search.startswith("•")
            assert not search.startswith("Search:")


class TestFallbackBehavior:
    """Tests for fallback to parent class when not enough searches."""

    def test_falls_back_to_parent_when_insufficient(self):
        """Test fallback to parent when not enough searches generated."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        # Return only one search but request 3
        mock_model.invoke.return_value = Mock(content="only one search")

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {
            "names": ["test"],
            "temporal": ["2024"],
            "descriptors": ["desc"],
        }

        with patch.object(
            FlexibleBrowseCompQuestionGenerator.__bases__[0],
            "_generate_progressive_searches",
            return_value=["parent search 1", "parent search 2"],
        ):
            searches = generator._generate_progressive_searches(
                query="Test",
                current_knowledge="",
                entities=entities,
                questions_by_iteration={},
                results_by_iteration={},
                num_questions=3,
                iteration=1,
            )

            # Should have combined results
            assert len(searches) == 3

    def test_no_fallback_when_sufficient(self):
        """Test no fallback when enough searches are generated."""
        from local_deep_research.advanced_search_system.questions.flexible_browsecomp_question import (
            FlexibleBrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="search one\nsearch two\nsearch three"
        )

        generator = FlexibleBrowseCompQuestionGenerator(mock_model)
        entities = {"names": [], "temporal": [], "descriptors": []}

        with patch.object(
            FlexibleBrowseCompQuestionGenerator.__bases__[0],
            "_generate_progressive_searches",
        ):
            searches = generator._generate_progressive_searches(
                query="Test",
                current_knowledge="",
                entities=entities,
                questions_by_iteration={},
                results_by_iteration={},
                num_questions=3,
                iteration=1,
            )

            # Should not call parent if sufficient
            # Note: This might still call parent in some edge cases
            assert len(searches) == 3
