"""
Comprehensive tests for PrecisionExtractionHandler.
Tests question type detection, various extraction methods, and edge cases.
"""

import pytest
from unittest.mock import Mock


class TestPrecisionExtractionHandlerInit:
    """Tests for PrecisionExtractionHandler initialization."""

    def test_init_with_llm_only(self, mock_llm):
        """Test initialization with just an LLM."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        assert handler.llm == mock_llm
        assert handler.settings_snapshot == {}

    def test_init_creates_answer_patterns(self, mock_llm):
        """Test initialization creates regex patterns."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        assert "full_name" in handler.answer_patterns
        assert "year" in handler.answer_patterns
        assert "number" in handler.answer_patterns
        assert "dimension" in handler.answer_patterns
        assert "score" in handler.answer_patterns

    def test_init_with_settings_snapshot(
        self, mock_llm, settings_with_fact_checking
    ):
        """Test initialization with settings snapshot."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(
            llm=mock_llm, settings_snapshot=settings_with_fact_checking
        )

        assert handler.settings_snapshot == settings_with_fact_checking


class TestPrecisionExtractionIdentifyQuestionType:
    """Tests for _identify_question_type method."""

    @pytest.fixture
    def handler(self, mock_llm):
        """Create a PrecisionExtractionHandler."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        return PrecisionExtractionHandler(llm=mock_llm)

    def test_identifies_full_name_question(self, handler):
        """Test detection of full name questions."""
        assert (
            handler._identify_question_type(
                "What is the full name of the founder?"
            )
            == "full_name"
        )

    def test_identifies_name_question_with_name_of(self, handler):
        """Test detection of name questions with 'name of'."""
        assert (
            handler._identify_question_type("What is the name of the CEO?")
            == "name"
        )

    def test_identifies_name_question_with_who_was(self, handler):
        """Test detection of name questions with 'who was'."""
        assert (
            handler._identify_question_type("Who was the inventor?") == "name"
        )

    def test_identifies_name_question_with_who_is(self, handler):
        """Test detection of name questions with 'who is'."""
        assert (
            handler._identify_question_type("Who is the current president?")
            == "name"
        )

    def test_identifies_location_question_with_where(self, handler):
        """Test detection of location questions with 'where'."""
        assert (
            handler._identify_question_type("Where did it happen?")
            == "location"
        )

    def test_identifies_location_question_with_city(self, handler):
        """Test detection of location questions mentioning city."""
        assert (
            handler._identify_question_type(
                "In which city was the company founded?"
            )
            == "location"
        )

    def test_identifies_location_question_with_country(self, handler):
        """Test detection of location questions mentioning country."""
        assert (
            handler._identify_question_type("What country is it located in?")
            == "location"
        )

    def test_identifies_temporal_question_with_when(self, handler):
        """Test detection of temporal questions with 'when'."""
        assert (
            handler._identify_question_type("When was the company founded?")
            == "temporal"
        )

    def test_identifies_temporal_question_with_year(self, handler):
        """Test detection of temporal questions with 'year'."""
        assert (
            handler._identify_question_type("In what year did the war begin?")
            == "temporal"
        )

    def test_identifies_temporal_question_with_date(self, handler):
        """Test detection of temporal questions with 'date'."""
        assert (
            handler._identify_question_type("What is the date of the event?")
            == "temporal"
        )

    def test_identifies_number_question_with_how_many(self, handler):
        """Test detection of number questions with 'how many'."""
        assert (
            handler._identify_question_type("How many employees work there?")
            == "number"
        )

    def test_identifies_number_question_with_how_much(self, handler):
        """Test detection of number questions with 'how much'."""
        assert (
            handler._identify_question_type("How much did it cost?") == "number"
        )

    def test_identifies_score_question(self, handler):
        """Test detection of score questions."""
        assert (
            handler._identify_question_type("What was the final score?")
            == "score"
        )

    def test_identifies_score_question_with_result(self, handler):
        """Test detection of score questions with 'result'."""
        assert (
            handler._identify_question_type("What was the result of the match?")
            == "score"
        )

    def test_identifies_dimension_question_with_height(self, handler):
        """Test detection of dimension questions with 'height'."""
        assert (
            handler._identify_question_type(
                "What is the height of the building?"
            )
            == "dimension"
        )

    def test_identifies_dimension_question_with_tall(self, handler):
        """Test detection of dimension questions with 'tall'."""
        assert (
            handler._identify_question_type("How tall is the tower?")
            == "dimension"
        )

    def test_identifies_dimension_question_with_length(self, handler):
        """Test detection of dimension questions with 'length'."""
        assert (
            handler._identify_question_type("What is the length of the bridge?")
            == "dimension"
        )

    def test_identifies_single_choice_question(self, handler):
        """Test detection of single choice questions."""
        assert (
            handler._identify_question_type("Which one of these is correct?")
            == "single_choice"
        )

    def test_returns_general_for_ambiguous_question(self, handler):
        """Test that ambiguous questions return 'general'."""
        assert (
            handler._identify_question_type("What happened in the meeting?")
            == "general"
        )

    def test_case_insensitive_detection(self, handler):
        """Test case-insensitive question type detection."""
        assert (
            handler._identify_question_type("WHAT IS THE FULL NAME?")
            == "full_name"
        )
        assert (
            handler._identify_question_type("WHERE DID IT HAPPEN?")
            == "location"
        )


class TestPrecisionExtractionRegexPatterns:
    """Tests for regex pattern matching."""

    @pytest.fixture
    def handler(self, mock_llm):
        """Create a PrecisionExtractionHandler."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        return PrecisionExtractionHandler(llm=mock_llm)

    def test_full_name_pattern_matches_two_words(self, handler):
        """Test full name pattern matches two-word names."""
        matches = handler.answer_patterns["full_name"].findall(
            "John Smith went to the store."
        )
        assert "John Smith" in matches

    def test_full_name_pattern_matches_three_words(self, handler):
        """Test full name pattern matches three-word names."""
        matches = handler.answer_patterns["full_name"].findall(
            "John Michael Smith was born."
        )
        assert "John Michael Smith" in matches

    def test_year_pattern_matches_1900s(self, handler):
        """Test year pattern matches 1900s years."""
        matches = handler.answer_patterns["year"].findall("Founded in 1995.")
        assert "1995" in matches

    def test_year_pattern_matches_2000s(self, handler):
        """Test year pattern matches 2000s years."""
        matches = handler.answer_patterns["year"].findall("Launched in 2024.")
        assert "2024" in matches

    def test_number_pattern_matches_integer(self, handler):
        """Test number pattern matches integers."""
        matches = handler.answer_patterns["number"].findall(
            "There are 42 employees."
        )
        assert "42" in matches

    def test_number_pattern_matches_decimal(self, handler):
        """Test number pattern matches decimals."""
        matches = handler.answer_patterns["number"].findall(
            "The cost is 19.99 dollars."
        )
        assert "19.99" in matches

    def test_dimension_pattern_matches_meters(self, handler):
        """Test dimension pattern matches meters."""
        matches = handler.answer_patterns["dimension"].findall(
            "The tower is 324 meters tall."
        )
        assert ("324", "meters") in matches

    def test_dimension_pattern_matches_feet(self, handler):
        """Test dimension pattern matches feet."""
        matches = handler.answer_patterns["dimension"].findall(
            "It is 100 feet high."
        )
        assert ("100", "feet") in matches

    def test_dimension_pattern_matches_kg(self, handler):
        """Test dimension pattern matches kilograms."""
        matches = handler.answer_patterns["dimension"].findall("Weighs 50 kg.")
        assert ("50", "kg") in matches

    def test_score_pattern_matches_hyphen_format(self, handler):
        """Test score pattern matches X-Y format."""
        matches = handler.answer_patterns["score"].findall(
            "Final score was 3-2."
        )
        assert ("3", "2") in matches

    def test_score_pattern_matches_dash_format(self, handler):
        """Test score pattern matches X–Y format (en-dash)."""
        matches = handler.answer_patterns["score"].findall(
            "Final score was 3–2."
        )
        assert ("3", "2") in matches


class TestPrecisionExtractionMethods:
    """Tests for individual extraction methods."""

    @pytest.fixture
    def handler(self, mock_llm):
        """Create a PrecisionExtractionHandler."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        return PrecisionExtractionHandler(llm=mock_llm)

    def test_extract_full_name_with_llm(self, handler, mock_llm):
        """Test _extract_full_name calls LLM."""
        mock_llm.invoke.return_value = Mock(
            content="Full name: John Michael Smith"
        )

        result = handler._extract_full_name(
            "John Smith was a scientist.", "Who is it?", "Sources"
        )

        mock_llm.invoke.assert_called_once()
        assert "John Michael Smith" in result

    def test_extract_full_name_fallback_to_longest(self, handler, mock_llm):
        """Test _extract_full_name falls back to longest name when LLM doesn't identify."""
        mock_llm.invoke.return_value = Mock(content="No full name identified")

        content = "John Michael Smith and John Smith were scientists."
        result = handler._extract_full_name(content, "Who is it?", "Sources")

        # Should pick the longer name
        assert "John Michael Smith" in result

    def test_extract_full_name_handles_exception(self, handler, mock_llm):
        """Test _extract_full_name handles exceptions gracefully."""
        mock_llm.invoke.side_effect = Exception("LLM error")

        result = handler._extract_full_name(
            "John Smith content", "Who?", "Sources"
        )

        # Should return original content
        assert "John Smith content" == result

    def test_extract_best_name_returns_most_frequent(self, handler):
        """Test _extract_best_name returns most frequent name."""
        content = "John Smith said hello. Smith was happy. John Smith left."

        result = handler._extract_best_name(content, "Who is it?", "Sources")

        assert "John Smith" in result

    def test_extract_dimension_with_llm(self, handler, mock_llm):
        """Test _extract_dimension calls LLM."""
        mock_llm.invoke.return_value = Mock(content="324 meters")

        result = handler._extract_dimension(
            "The tower is 324 meters tall.", "How tall?", "Sources"
        )

        mock_llm.invoke.assert_called_once()
        assert "324" in result

    def test_extract_dimension_fallback_to_pattern(self, handler, mock_llm):
        """Test _extract_dimension falls back to regex pattern."""
        mock_llm.invoke.return_value = Mock(content="No clear answer")

        content = "The tower is 324 meters tall."
        result = handler._extract_dimension(content, "How tall?", "Sources")

        # Should still extract the measurement
        assert "324" in result or "meters" in result or content == result

    def test_extract_score_with_multiple_scores(self, handler, mock_llm):
        """Test _extract_score when multiple scores present."""
        mock_llm.invoke.return_value = Mock(content="3-2")

        content = "Halftime was 1-1. Final score was 3-2."
        result = handler._extract_score(
            content, "What was the final score?", "Sources"
        )

        mock_llm.invoke.assert_called_once()
        assert "3-2" in result

    def test_extract_score_no_scores_returns_content(self, handler, mock_llm):
        """Test _extract_score returns content when no scores found."""
        content = "The game was exciting."
        result = handler._extract_score(
            content, "What was the score?", "Sources"
        )

        assert result == content

    def test_extract_temporal_with_llm(self, handler, mock_llm):
        """Test _extract_temporal calls LLM."""
        mock_llm.invoke.return_value = Mock(content="1998")

        content = "Founded in 1998, acquired in 2015."
        result = handler._extract_temporal(
            content, "When was it founded?", "Sources"
        )

        mock_llm.invoke.assert_called_once()
        assert "1998" in result

    def test_extract_temporal_no_years_returns_content(self, handler, mock_llm):
        """Test _extract_temporal returns content when no years found."""
        content = "It happened long ago."
        result = handler._extract_temporal(content, "When?", "Sources")

        assert result == content

    def test_extract_number_with_llm(self, handler, mock_llm):
        """Test _extract_number calls LLM."""
        mock_llm.invoke.return_value = Mock(content="42")

        content = "There are 42 employees and 10 managers."
        result = handler._extract_number(
            content, "How many employees?", "Sources"
        )

        mock_llm.invoke.assert_called_once()
        assert "42" in result

    def test_extract_single_answer_removes_alternatives(
        self, handler, mock_llm
    ):
        """Test _extract_single_answer removes comma/and/or alternatives."""
        mock_llm.invoke.return_value = Mock(
            content="Option A, Option B, and Option C"
        )

        result = handler._extract_single_answer(
            "Multiple options content", "Which one?", "Sources"
        )

        # Should only have first option
        assert "Option A" in result
        assert "Option B" not in result.split(".")[0]

    def test_extract_key_facts_from_previous_knowledge(self, handler, mock_llm):
        """Test _extract_key_facts extracts from previous knowledge."""
        mock_llm.invoke.return_value = Mock(
            content="Key facts: Founded 1998, Location: NYC"
        )

        result = handler._extract_key_facts(
            "Previous knowledge about the company", "name"
        )

        mock_llm.invoke.assert_called_once()
        assert len(result) <= 500

    def test_extract_key_facts_handles_exception(self, handler, mock_llm):
        """Test _extract_key_facts handles exception gracefully."""
        mock_llm.invoke.side_effect = Exception("LLM error")

        result = handler._extract_key_facts("Previous knowledge", "name")

        # Should return truncated previous knowledge
        assert "Previous knowledge" in result


class TestPrecisionExtractionAnalyzeInitial:
    """Tests for analyze_initial method."""

    def test_analyze_initial_returns_content_and_documents(
        self, mock_llm, sample_search_results
    ):
        """Test analyze_initial returns proper structure."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        result = handler.analyze_initial(
            "What is the topic?", sample_search_results
        )

        assert "content" in result
        assert "documents" in result
        assert len(result["documents"]) == 3

    def test_analyze_initial_prompt_includes_question_type(
        self, mock_llm, sample_search_results
    ):
        """Test analyze_initial prompt includes detected question type."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        mock_llm.invoke.return_value = Mock(
            content="The full name is John Smith. [1]"
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        handler.analyze_initial(
            "What is the full name of the founder?", sample_search_results
        )

        call_args = mock_llm.invoke.call_args_list[0][0][0]
        assert "full_name" in call_args

    def test_analyze_initial_prompt_includes_precision_instructions(
        self, mock_llm, sample_search_results
    ):
        """Test analyze_initial prompt includes precision extraction instructions."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        mock_llm.invoke.return_value = Mock(content="Answer [1]")

        handler = PrecisionExtractionHandler(llm=mock_llm)

        handler.analyze_initial("What is it?", sample_search_results)

        call_args = mock_llm.invoke.call_args_list[0][0][0]
        assert "PRECISION" in call_args
        assert "EXACT answer" in call_args

    def test_analyze_initial_applies_extraction_for_full_name(
        self, mock_llm, name_search_results
    ):
        """Test analyze_initial applies extraction for full_name questions."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        # First call returns LLM response, second call extracts name
        mock_llm.invoke.side_effect = [
            Mock(content="The person is John Smith."),
            Mock(content="Full name: John Michael William Smith"),
        ]

        handler = PrecisionExtractionHandler(llm=mock_llm)

        handler.analyze_initial("What is the full name?", name_search_results)

        # Should have called extraction
        assert mock_llm.invoke.call_count >= 1

    def test_analyze_initial_handles_string_response(
        self, mock_llm_string_response, sample_search_results
    ):
        """Test analyze_initial handles string LLM response."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm_string_response)

        result = handler.analyze_initial("What is it?", sample_search_results)

        assert "content" in result


class TestPrecisionExtractionAnalyzeFollowup:
    """Tests for analyze_followup method."""

    def test_analyze_followup_returns_content_and_documents(
        self, mock_llm, sample_search_results, sample_previous_knowledge
    ):
        """Test analyze_followup returns proper structure."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        result = handler.analyze_followup(
            "Follow-up question?",
            sample_search_results,
            sample_previous_knowledge,
            nr_of_links=3,
        )

        assert "content" in result
        assert "documents" in result

    def test_analyze_followup_extracts_key_facts(
        self, mock_llm, sample_search_results, sample_previous_knowledge
    ):
        """Test analyze_followup extracts key facts from previous knowledge."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        mock_llm.invoke.return_value = Mock(content="Key facts extracted")

        handler = PrecisionExtractionHandler(llm=mock_llm)

        handler.analyze_followup(
            "What year was it founded?",
            sample_search_results,
            sample_previous_knowledge,
            nr_of_links=0,
        )

        # Should have called LLM for key facts extraction
        first_call = mock_llm.invoke.call_args_list[0][0][0]
        assert (
            "key facts" in first_call.lower() or "extract" in first_call.lower()
        )

    def test_analyze_followup_applies_nr_of_links_offset(
        self, mock_llm, sample_search_results, sample_previous_knowledge
    ):
        """Test analyze_followup applies nr_of_links offset to document indices."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        result = handler.analyze_followup(
            "Question?",
            sample_search_results,
            sample_previous_knowledge,
            nr_of_links=10,
        )

        # First document should have index 11 (10 + 1)
        assert result["documents"][0].metadata["index"] == 11


class TestPrecisionExtractionEdgeCases:
    """Tests for edge cases and error handling."""

    def test_empty_search_results(self, mock_llm, empty_search_results):
        """Test handling of empty search results."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        result = handler.analyze_initial("What is it?", empty_search_results)

        assert result["documents"] == []

    def test_string_search_results(self, mock_llm, string_search_results):
        """Test handling of string search results."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        result = handler.analyze_initial("What is it?", string_search_results)

        assert result["documents"] == []

    def test_with_output_instructions(
        self, mock_llm, sample_search_results, settings_with_output_instructions
    ):
        """Test that output instructions are included in prompts."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        mock_llm.invoke.return_value = Mock(content="Response")

        handler = PrecisionExtractionHandler(
            llm=mock_llm, settings_snapshot=settings_with_output_instructions
        )

        handler.analyze_initial("What is it?", sample_search_results)

        call_args = mock_llm.invoke.call_args_list[0][0][0]
        assert "formal academic English" in call_args

    def test_general_question_type_no_extraction(
        self, mock_llm, sample_search_results
    ):
        """Test that general question type doesn't apply special extraction."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        mock_llm.invoke.return_value = Mock(content="General response")

        handler = PrecisionExtractionHandler(llm=mock_llm)

        result = handler.analyze_initial(
            "What happened?", sample_search_results
        )

        # Should just return the LLM response without extraction modifications
        assert result["content"] == "General response"

    def test_apply_precision_extraction_returns_content_for_unknown_type(
        self, mock_llm
    ):
        """Test _apply_precision_extraction returns content for unknown question type."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        handler = PrecisionExtractionHandler(llm=mock_llm)

        result = handler._apply_precision_extraction(
            "Original content", "Query?", "unknown_type", "Sources"
        )

        assert result == "Original content"


class TestPrecisionExtractionDimensionContextAwareness:
    """Tests for dimension extraction context awareness."""

    @pytest.fixture
    def handler(self, mock_llm):
        """Create a PrecisionExtractionHandler."""
        from local_deep_research.citation_handlers.precision_extraction_handler import (
            PrecisionExtractionHandler,
        )

        return PrecisionExtractionHandler(llm=mock_llm)

    def test_dimension_types_include_height(self, handler):
        """Test dimension type detection includes height keywords."""
        question_type = handler._identify_question_type(
            "What is the height of the building?"
        )
        assert question_type == "dimension"

    def test_dimension_types_include_size(self, handler):
        """Test dimension type detection includes size keywords."""
        question_type = handler._identify_question_type(
            "What is the size of the building?"
        )
        assert question_type == "dimension"

    def test_dimension_types_include_length(self, handler):
        """Test dimension type detection includes length keywords."""
        question_type = handler._identify_question_type(
            "What is the length of the bridge?"
        )
        assert question_type == "dimension"

    def test_dimension_types_include_width(self, handler):
        """Test dimension type detection includes width keywords."""
        question_type = handler._identify_question_type(
            "How wide is the river?"
        )
        assert question_type == "dimension"
