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

Tests cover:
- BrowseCompQuestionGenerator initialization
- _extract_entities method
- _expand_temporal_ranges method
- _generate_initial_searches method
- _generate_progressive_searches method
- _format_previous_searches method
- _was_searched method
- generate_questions method
"""

from unittest.mock import Mock


class TestBrowseCompQuestionGeneratorInit:
    """Tests for BrowseCompQuestionGenerator initialization."""

    def test_default_initialization(self):
        """Test initialization with default parameters."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        generator = BrowseCompQuestionGenerator(mock_model)

        assert generator.model is mock_model
        assert generator.extracted_entities == {}
        assert generator.search_progression == []
        assert generator.knowledge_truncate_length == 1500
        assert generator.previous_searches_limit == 10

    def test_custom_initialization(self):
        """Test initialization with custom parameters."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(
            Mock(),
            knowledge_truncate_length=2000,
            previous_searches_limit=15,
        )

        assert generator.knowledge_truncate_length == 2000
        assert generator.previous_searches_limit == 15

    def test_inherits_from_base_question_generator(self):
        """Test that it inherits from BaseQuestionGenerator."""
        from local_deep_research.advanced_search_system.questions.base_question import (
            BaseQuestionGenerator,
        )
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        assert issubclass(BrowseCompQuestionGenerator, BaseQuestionGenerator)


class TestExtractEntities:
    """Tests for _extract_entities method."""

    def test_extracts_temporal_entities(self):
        """Test extracting temporal entities from query."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="TEMPORAL: 2018, 2019\nNUMERICAL:\nNAMES:\nLOCATIONS:\nDESCRIPTORS:"
        )

        generator = BrowseCompQuestionGenerator(mock_model)
        entities = generator._extract_entities("Events in 2018 and 2019")

        assert "2018" in entities["temporal"]
        assert "2019" in entities["temporal"]

    def test_extracts_names(self):
        """Test extracting name entities."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="TEMPORAL:\nNUMERICAL:\nNAMES: Einstein, Newton\nLOCATIONS:\nDESCRIPTORS:"
        )

        generator = BrowseCompQuestionGenerator(mock_model)
        entities = generator._extract_entities("Physicists Einstein and Newton")

        assert "Einstein" in entities["names"]
        assert "Newton" in entities["names"]

    def test_extracts_locations(self):
        """Test extracting location entities."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="TEMPORAL:\nNUMERICAL:\nNAMES:\nLOCATIONS: Paris, London\nDESCRIPTORS:"
        )

        generator = BrowseCompQuestionGenerator(mock_model)
        entities = generator._extract_entities("Cities Paris and London")

        assert "Paris" in entities["locations"]
        assert "London" in entities["locations"]

    def test_extracts_descriptors(self):
        """Test extracting descriptor entities."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="TEMPORAL:\nNUMERICAL:\nNAMES:\nLOCATIONS:\nDESCRIPTORS: famous, historical"
        )

        generator = BrowseCompQuestionGenerator(mock_model)
        entities = generator._extract_entities("Famous historical events")

        assert "famous" in entities["descriptors"]
        assert "historical" in entities["descriptors"]

    def test_handles_response_without_content_attr(self):
        """Test handling response that is a string."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        # Return string instead of object with content attribute
        mock_model.invoke.return_value = (
            "TEMPORAL: 2020\nNUMERICAL:\nNAMES:\nLOCATIONS:\nDESCRIPTORS:"
        )

        generator = BrowseCompQuestionGenerator(mock_model)
        entities = generator._extract_entities("Test query")

        # Should handle gracefully
        assert isinstance(entities, dict)

    def test_extracts_numerical_entities(self):
        """Test extracting numerical entities."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="TEMPORAL:\nNUMERICAL: 42, 100\nNAMES:\nLOCATIONS:\nDESCRIPTORS:"
        )

        generator = BrowseCompQuestionGenerator(mock_model)
        entities = generator._extract_entities("Numbers 42 and 100")

        assert "42" in entities["numerical"]
        assert "100" in entities["numerical"]


class TestExpandTemporalRanges:
    """Tests for _expand_temporal_ranges method."""

    def test_expands_year_range_with_dash(self):
        """Test expanding year ranges with dash separator."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        result = generator._expand_temporal_ranges(["2018-2020"])

        assert "2018" in result
        assert "2019" in result
        assert "2020" in result

    def test_expands_year_range_with_to(self):
        """Test expanding year ranges with 'to' separator."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        result = generator._expand_temporal_ranges(["2015 to 2017"])

        assert "2015" in result
        assert "2016" in result
        assert "2017" in result

    def test_expands_year_range_with_and(self):
        """Test expanding year ranges with 'and' separator."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        result = generator._expand_temporal_ranges(["between 2010 and 2012"])

        assert "2010" in result
        assert "2011" in result
        assert "2012" in result

    def test_extracts_single_year(self):
        """Test extracting single year from text."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        result = generator._expand_temporal_ranges(["in 2023"])

        assert "2023" in result

    def test_removes_duplicates(self):
        """Test that duplicates are removed."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        result = generator._expand_temporal_ranges(["2020", "2020-2021"])

        # 2020 appears in both but should only be in result once
        assert result.count("2020") == 1

    def test_handles_non_year_temporal(self):
        """Test handling non-year temporal entities."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        result = generator._expand_temporal_ranges(["summer", "morning"])

        assert "summer" in result
        assert "morning" in result


class TestGenerateInitialSearches:
    """Tests for _generate_initial_searches method."""

    def test_always_includes_original_query(self):
        """Test that original query is always included."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        entities = {
            "temporal": [],
            "numerical": [],
            "names": [],
            "locations": [],
            "descriptors": [],
        }

        searches = generator._generate_initial_searches(
            "Original query", entities, 3
        )

        assert "Original query" in searches

    def test_returns_only_original_when_one_requested(self):
        """Test that only original query is returned when 1 question requested."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        entities = {
            "temporal": ["2020"],
            "numerical": [],
            "names": ["Einstein"],
            "locations": [],
            "descriptors": [],
        }

        searches = generator._generate_initial_searches(
            "Test query", entities, 1
        )

        assert len(searches) == 1
        assert searches[0] == "Test query"

    def test_includes_name_searches(self):
        """Test that name-based searches are included."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        entities = {
            "temporal": [],
            "numerical": [],
            "names": ["Einstein"],
            "locations": [],
            "descriptors": ["physicist"],
        }

        searches = generator._generate_initial_searches(
            "Test query", entities, 5
        )

        assert any("Einstein" in s for s in searches)

    def test_includes_temporal_searches(self):
        """Test that year-based searches are included."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        entities = {
            "temporal": ["2018", "2019"],
            "numerical": [],
            "names": ["Einstein"],
            "locations": [],
            "descriptors": [],
        }

        searches = generator._generate_initial_searches(
            "Test query", entities, 5
        )

        assert any("2018" in s or "2019" in s for s in searches)

    def test_includes_location_searches(self):
        """Test that location-based searches are included."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        entities = {
            "temporal": [],
            "numerical": [],
            "names": [],
            "locations": ["Paris"],
            "descriptors": ["historic"],
        }

        searches = generator._generate_initial_searches(
            "Test query", entities, 5
        )

        assert any("Paris" in s for s in searches)

    def test_removes_duplicate_searches(self):
        """Test that duplicate searches are removed."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        entities = {
            "temporal": [],
            "numerical": [],
            "names": ["test"],
            "locations": ["test"],  # Same as name
            "descriptors": [],
        }

        searches = generator._generate_initial_searches("Test", entities, 5)

        # Case-insensitive duplicate check
        lower_searches = [s.lower() for s in searches]
        assert len(lower_searches) == len(set(lower_searches))

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

        generator = BrowseCompQuestionGenerator(Mock())
        entities = {
            "temporal": ["2018", "2019", "2020"],
            "numerical": ["100"],
            "names": ["A", "B", "C"],
            "locations": ["X", "Y", "Z"],
            "descriptors": ["a", "b", "c"],
        }

        searches = generator._generate_initial_searches(
            "Test query", entities, 3
        )

        assert len(searches) <= 3


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

    def test_generates_from_model_response(self):
        """Test generating searches from model response."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

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

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

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

        assert len(searches) == 3

    def test_detects_failing_searches(self):
        """Test detection of failing searches (0 results)."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

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

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

        # Simulate 3+ iterations with 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,
        )

        # Verify the prompt mentions failing searches
        call_args = mock_model.invoke.call_args[0][0]
        assert "0 results" in call_args or "NARROW" in call_args

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

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

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

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

        for search in searches:
            assert not search.startswith("Q:")
            assert not search.startswith("-")
            assert not search.startswith("*")
            assert not search.startswith("•")

    def test_adds_year_searches_when_needed(self):
        """Test that year-based searches are added programmatically."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(content="only one")

        generator = BrowseCompQuestionGenerator(mock_model)
        entities = {
            "temporal": ["2020", "2021"],
            "numerical": [],
            "names": ["Einstein"],
            "locations": [],
            "descriptors": [],
        }

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

        # Should have generated year-based searches to fill quota
        assert any("2020" in s or "2021" in s for s in searches)

    def test_truncates_knowledge(self):
        """Test that knowledge is truncated when configured."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

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

        generator = BrowseCompQuestionGenerator(
            mock_model, knowledge_truncate_length=50
        )
        entities = {
            "temporal": [],
            "numerical": [],
            "names": [],
            "locations": [],
            "descriptors": [],
        }

        long_knowledge = "A" * 200

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

        call_args = mock_model.invoke.call_args[0][0]
        # Knowledge should be truncated in prompt
        assert call_args.count("A") <= 50


class TestFormatPreviousSearches:
    """Tests for _format_previous_searches method."""

    def test_formats_with_result_counts(self):
        """Test formatting with result counts."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        questions_by_iteration = {
            1: ["query 1", "query 2"],
            2: ["query 3"],
        }
        results_by_iteration = {1: 5, 2: 10}

        formatted = generator._format_previous_searches(
            questions_by_iteration, results_by_iteration
        )

        assert "5 results" in formatted
        assert "10 results" in formatted

    def test_handles_missing_result_counts(self):
        """Test handling missing result counts."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        questions_by_iteration = {1: ["query 1"]}
        results_by_iteration = {}  # No results data

        formatted = generator._format_previous_searches(
            questions_by_iteration, results_by_iteration
        )

        # Should not include "? results"
        assert "? results" not in formatted

    def test_respects_previous_searches_limit(self):
        """Test that previous searches are limited."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(
            Mock(), previous_searches_limit=3
        )
        questions_by_iteration = {
            1: ["q1", "q2"],
            2: ["q3", "q4"],
            3: ["q5", "q6"],
        }
        results_by_iteration = {1: 1, 2: 2, 3: 3}

        formatted = generator._format_previous_searches(
            questions_by_iteration, results_by_iteration
        )

        # Should only show last 3
        lines = formatted.strip().split("\n")
        assert len(lines) <= 3


class TestWasSearched:
    """Tests for _was_searched method."""

    def test_returns_true_for_searched_term(self):
        """Test that True is returned for already searched term."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        questions_by_iteration = {
            1: ["Einstein physics 2020"],
        }

        assert generator._was_searched("Einstein", questions_by_iteration)

    def test_returns_false_for_new_term(self):
        """Test that False is returned for new term."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        questions_by_iteration = {
            1: ["Einstein physics 2020"],
        }

        assert not generator._was_searched("Newton", questions_by_iteration)

    def test_case_insensitive_matching(self):
        """Test that matching is case-insensitive."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        questions_by_iteration = {
            1: ["EINSTEIN physics"],
        }

        assert generator._was_searched("einstein", questions_by_iteration)

    def test_handles_non_list_values(self):
        """Test handling of non-list values in questions_by_iteration."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        generator = BrowseCompQuestionGenerator(Mock())
        questions_by_iteration = {
            1: "not a list",  # Invalid but should handle gracefully
        }

        # Should not raise, just return False
        assert not generator._was_searched("test", questions_by_iteration)


class TestGenerateQuestions:
    """Tests for generate_questions method."""

    def test_first_iteration_extracts_entities(self):
        """Test that first iteration extracts entities."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="TEMPORAL: 2020\nNUMERICAL:\nNAMES: Test\nLOCATIONS:\nDESCRIPTORS:"
        )

        generator = BrowseCompQuestionGenerator(mock_model)
        questions = generator.generate_questions(
            current_knowledge="",
            query="Test query",
            questions_per_iteration=3,
            iteration=1,
        )

        assert generator.extracted_entities  # Should have extracted entities
        assert len(questions) <= 3

    def test_subsequent_iteration_uses_progressive_searches(self):
        """Test that subsequent iterations use progressive searches."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

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

        generator = BrowseCompQuestionGenerator(mock_model)
        # Pre-populate entities
        generator.extracted_entities = {
            "temporal": ["2020"],
            "numerical": [],
            "names": ["Test"],
            "locations": [],
            "descriptors": [],
        }

        questions = generator.generate_questions(
            current_knowledge="Some knowledge",
            query="Test query",
            questions_per_iteration=3,
            iteration=2,
        )

        assert len(questions) == 3

    def test_handles_empty_questions_by_iteration(self):
        """Test handling of None questions_by_iteration."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

        mock_model = Mock()
        mock_model.invoke.return_value = Mock(
            content="TEMPORAL:\nNUMERICAL:\nNAMES:\nLOCATIONS:\nDESCRIPTORS:"
        )

        generator = BrowseCompQuestionGenerator(mock_model)
        questions = generator.generate_questions(
            current_knowledge="",
            query="Test",
            questions_per_iteration=1,
            questions_by_iteration=None,
            iteration=1,
        )

        assert isinstance(questions, list)

    def test_handles_empty_results_by_iteration(self):
        """Test handling of None results_by_iteration."""
        from local_deep_research.advanced_search_system.questions.browsecomp_question import (
            BrowseCompQuestionGenerator,
        )

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

        generator = BrowseCompQuestionGenerator(mock_model)
        generator.extracted_entities = {
            "temporal": [],
            "numerical": [],
            "names": [],
            "locations": [],
            "descriptors": [],
        }

        questions = generator.generate_questions(
            current_knowledge="",
            query="Test",
            questions_per_iteration=1,
            results_by_iteration=None,
            iteration=2,
        )

        assert isinstance(questions, list)
