"""
Tests for benchmarks/graders.py

Tests cover:
- extract_answer_from_response function
- grade_single_result with mocked LLM
- get_evaluation_llm configuration
"""

from unittest.mock import Mock, patch, MagicMock


class TestExtractAnswerFromResponse:
    """Tests for the extract_answer_from_response function."""

    def test_browsecomp_extracts_exact_answer(self):
        """Test extraction of exact answer from BrowseComp response."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = """
Based on my research, I found the following information.

Exact Answer: 42
Confidence: 95%
"""
        result = extract_answer_from_response(
            response, dataset_type="browsecomp"
        )

        assert result["extracted_answer"] == "42"
        assert result["confidence"] == "95"

    def test_browsecomp_missing_answer_returns_none(self):
        """Test handling of missing answer in BrowseComp response."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = "Some response without the expected format"
        result = extract_answer_from_response(
            response, dataset_type="browsecomp"
        )

        assert result["extracted_answer"] == "None"
        assert result["confidence"] == "100"

    def test_browsecomp_missing_confidence_defaults(self):
        """Test that missing confidence defaults to 100."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = "Exact Answer: Paris"
        result = extract_answer_from_response(
            response, dataset_type="browsecomp"
        )

        assert result["extracted_answer"] == "Paris"
        assert result["confidence"] == "100"

    def test_simpleqa_returns_full_response(self):
        """Test that SimpleQA returns the full response as the answer."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = "The capital of France is Paris."
        result = extract_answer_from_response(response, dataset_type="simpleqa")

        assert result["extracted_answer"] == response
        assert result["confidence"] == "100"

    def test_removes_citations_from_response(self):
        """Test that citations are removed from the response."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = "The answer is 42 [1] according to the source [2][3]."
        result = extract_answer_from_response(response, dataset_type="simpleqa")

        assert "[1]" not in result["extracted_answer"]
        assert "[2]" not in result["extracted_answer"]
        assert "[3]" not in result["extracted_answer"]
        assert "42" in result["extracted_answer"]

    def test_browsecomp_case_insensitive(self):
        """Test that dataset type matching is case insensitive."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = "Exact Answer: test\nConfidence: 80%"
        result = extract_answer_from_response(
            response, dataset_type="BROWSECOMP"
        )

        assert result["extracted_answer"] == "test"
        assert result["confidence"] == "80"


class TestGetEvaluationLLM:
    """Tests for the get_evaluation_llm function."""

    @patch("local_deep_research.benchmarks.graders.get_llm")
    def test_uses_default_config(self, mock_get_llm):
        """Test that default config is used."""
        from local_deep_research.benchmarks.graders import (
            get_evaluation_llm,
            DEFAULT_EVALUATION_CONFIG,
        )

        mock_get_llm.return_value = Mock()

        get_evaluation_llm()

        mock_get_llm.assert_called_once()
        call_kwargs = mock_get_llm.call_args[1]
        assert (
            call_kwargs["model_name"] == DEFAULT_EVALUATION_CONFIG["model_name"]
        )
        assert (
            call_kwargs["temperature"]
            == DEFAULT_EVALUATION_CONFIG["temperature"]
        )

    @patch("local_deep_research.benchmarks.graders.get_llm")
    def test_custom_config_overrides_defaults(self, mock_get_llm):
        """Test that custom config overrides defaults."""
        from local_deep_research.benchmarks.graders import get_evaluation_llm

        mock_get_llm.return_value = Mock()

        custom_config = {
            "model_name": "custom-model",
            "temperature": 0.5,
        }
        get_evaluation_llm(custom_config=custom_config)

        call_kwargs = mock_get_llm.call_args[1]
        assert call_kwargs["model_name"] == "custom-model"
        assert call_kwargs["temperature"] == 0.5

    @patch("local_deep_research.benchmarks.graders.get_llm")
    def test_filters_unsupported_params(self, mock_get_llm):
        """Test that unsupported parameters are filtered out."""
        from local_deep_research.benchmarks.graders import get_evaluation_llm

        mock_get_llm.return_value = Mock()

        custom_config = {
            "model_name": "test-model",
            "unsupported_param": "value",
            "max_tokens": 1000,  # Not supported by LDR's get_llm
        }
        get_evaluation_llm(custom_config=custom_config)

        call_kwargs = mock_get_llm.call_args[1]
        assert "unsupported_param" not in call_kwargs
        assert "max_tokens" not in call_kwargs

    @patch("local_deep_research.benchmarks.graders.get_llm")
    def test_extracts_api_key_from_settings_snapshot(self, mock_get_llm):
        """Test that API key is extracted from settings snapshot."""
        from local_deep_research.benchmarks.graders import get_evaluation_llm

        mock_get_llm.return_value = Mock()

        settings_snapshot = {
            "llm.openai_endpoint.api_key": {"value": "test-api-key"}
        }
        get_evaluation_llm(settings_snapshot=settings_snapshot)

        # Should not raise and should call get_llm
        mock_get_llm.assert_called_once()


class TestGradeSingleResult:
    """Tests for the grade_single_result function."""

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grades_correctly(self, mock_get_eval_llm):
        """Test that grade_single_result grades correctly."""
        from local_deep_research.benchmarks.graders import grade_single_result

        # Mock the LLM response
        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = """
Extracted Answer: Paris
Reasoning: The model correctly identified Paris as the capital of France.
Correct: yes
"""
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        result_data = {
            "problem": "What is the capital of France?",
            "correct_answer": "Paris",
            "response": "The capital of France is Paris.",
        }

        graded = grade_single_result(result_data, dataset_type="simpleqa")

        assert graded["is_correct"] is True
        assert "Paris" in graded["extracted_by_grader"]

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_handles_grading_error(self, mock_get_eval_llm):
        """Test that grade_single_result handles errors gracefully."""
        from local_deep_research.benchmarks.graders import grade_single_result

        # Mock the LLM to raise an error during invoke (inside the try block)
        mock_llm = MagicMock()
        mock_llm.invoke.side_effect = Exception("LLM invoke error")
        mock_get_eval_llm.return_value = mock_llm

        result_data = {
            "problem": "test",
            "correct_answer": "answer",
            "response": "response",
        }

        graded = grade_single_result(result_data)

        assert graded["is_correct"] is False
        assert "grading_error" in graded

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_browsecomp_grading_format(self, mock_get_eval_llm):
        """Test BrowseComp-specific grading format extraction."""
        from local_deep_research.benchmarks.graders import grade_single_result

        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = """
extracted_final_answer: 42
reasoning: The model found the correct answer by analyzing the data.
correct: yes
confidence: 95
"""
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        result_data = {
            "problem": "What is the answer?",
            "correct_answer": "42",
            "response": "The answer is 42.",
        }

        graded = grade_single_result(result_data, dataset_type="browsecomp")

        assert graded["is_correct"] is True
        assert graded["extracted_by_grader"] == "42"
        assert graded["graded_confidence"] == "95"

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grading_with_no_judgment(self, mock_get_eval_llm):
        """Test grading when LLM doesn't provide clear judgment."""
        from local_deep_research.benchmarks.graders import grade_single_result

        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "Some response without proper format"
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        result_data = {
            "problem": "test",
            "correct_answer": "answer",
            "response": "response",
        }

        graded = grade_single_result(result_data, dataset_type="simpleqa")

        # Should default to False when no clear judgment
        assert graded["is_correct"] is False
        assert graded["extracted_by_grader"] == "None"


class TestGradeResults:
    """Tests for grade_results function (batch grading)."""

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_results_processes_all_items(self, mock_get_eval_llm):
        """Test that grade_results processes all items in file."""
        import tempfile
        import json
        from local_deep_research.benchmarks.graders import grade_results

        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = """
Extracted Answer: test
Reasoning: Test reasoning
Correct: yes
"""
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        with tempfile.TemporaryDirectory() as tmpdir:
            # Create input file
            input_file = f"{tmpdir}/input.jsonl"
            with open(input_file, "w") as f:
                for i in range(3):
                    f.write(
                        json.dumps(
                            {
                                "problem": f"Question {i}",
                                "correct_answer": f"Answer {i}",
                                "response": f"Response {i}",
                            }
                        )
                        + "\n"
                    )

            output_file = f"{tmpdir}/output.jsonl"

            results = grade_results(input_file, output_file)

            assert len(results) == 3
            assert all(r["is_correct"] for r in results)

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_results_invokes_progress_callback(self, mock_get_eval_llm):
        """Test that progress callback is invoked during grading."""
        import tempfile
        import json
        from local_deep_research.benchmarks.graders import grade_results

        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = (
            "Extracted Answer: test\nReasoning: test\nCorrect: yes"
        )
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        callback_invocations = []

        def progress_callback(idx, total, data):
            callback_invocations.append(
                {"idx": idx, "total": total, "data": data}
            )

        with tempfile.TemporaryDirectory() as tmpdir:
            input_file = f"{tmpdir}/input.jsonl"
            with open(input_file, "w") as f:
                f.write(
                    json.dumps(
                        {
                            "problem": "Q",
                            "correct_answer": "A",
                            "response": "R",
                        }
                    )
                    + "\n"
                )

            output_file = f"{tmpdir}/output.jsonl"

            grade_results(
                input_file, output_file, progress_callback=progress_callback
            )

            # Should have multiple invocations (grading and graded)
            assert len(callback_invocations) >= 2

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_results_handles_errors_gracefully(self, mock_get_eval_llm):
        """Test that grade_results handles individual grading errors."""
        import tempfile
        import json
        from local_deep_research.benchmarks.graders import grade_results

        mock_llm = MagicMock()
        # First call succeeds, second fails
        mock_response = MagicMock()
        mock_response.content = "Extracted Answer: test\nCorrect: yes"
        mock_llm.invoke.side_effect = [
            mock_response,
            Exception("Grading error"),
        ]
        mock_get_eval_llm.return_value = mock_llm

        with tempfile.TemporaryDirectory() as tmpdir:
            input_file = f"{tmpdir}/input.jsonl"
            with open(input_file, "w") as f:
                for i in range(2):
                    f.write(
                        json.dumps(
                            {
                                "problem": f"Q{i}",
                                "correct_answer": f"A{i}",
                                "response": f"R{i}",
                            }
                        )
                        + "\n"
                    )

            output_file = f"{tmpdir}/output.jsonl"

            results = grade_results(input_file, output_file)

            # Should have both results (one success, one error)
            assert len(results) == 2
            # First should be correct
            assert results[0]["is_correct"] is True
            # Second should have error
            assert "grading_error" in results[1]

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_results_writes_output_file(self, mock_get_eval_llm):
        """Test that grade_results writes to output file."""
        import tempfile
        import json
        from local_deep_research.benchmarks.graders import grade_results

        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "Extracted Answer: test\nCorrect: yes"
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        with tempfile.TemporaryDirectory() as tmpdir:
            input_file = f"{tmpdir}/input.jsonl"
            with open(input_file, "w") as f:
                f.write(
                    json.dumps(
                        {"problem": "Q", "correct_answer": "A", "response": "R"}
                    )
                    + "\n"
                )

            output_file = f"{tmpdir}/output.jsonl"

            grade_results(input_file, output_file)

            # Output file should exist
            with open(output_file, "r") as f:
                lines = f.readlines()

            assert len(lines) == 1
            result = json.loads(lines[0])
            assert "is_correct" in result


class TestHumanEvaluation:
    """Tests for human_evaluation function."""

    def test_human_evaluation_noninteractive_mode(self):
        """Test human evaluation in non-interactive mode."""
        import tempfile
        import json
        from local_deep_research.benchmarks.graders import human_evaluation

        with tempfile.TemporaryDirectory() as tmpdir:
            input_file = f"{tmpdir}/input.jsonl"
            with open(input_file, "w") as f:
                f.write(
                    json.dumps(
                        {
                            "problem": "What is 2+2?",
                            "correct_answer": "4",
                            "response": "The answer is 4.",
                            "extracted_answer": "4",
                        }
                    )
                    + "\n"
                )

            output_file = f"{tmpdir}/output.jsonl"

            results = human_evaluation(
                input_file, output_file, interactive=False
            )

            assert len(results) == 1
            # Non-interactive defaults to is_correct=False
            assert results[0]["is_correct"] is False
            assert results[0]["human_evaluation"] is True

    def test_human_evaluation_writes_output(self):
        """Test that human evaluation writes to output file."""
        import tempfile
        import json
        from local_deep_research.benchmarks.graders import human_evaluation

        with tempfile.TemporaryDirectory() as tmpdir:
            input_file = f"{tmpdir}/input.jsonl"
            with open(input_file, "w") as f:
                f.write(
                    json.dumps(
                        {
                            "problem": "Q",
                            "correct_answer": "A",
                            "response": "R",
                        }
                    )
                    + "\n"
                )

            output_file = f"{tmpdir}/output.jsonl"

            human_evaluation(input_file, output_file, interactive=False)

            with open(output_file, "r") as f:
                lines = f.readlines()

            assert len(lines) == 1
            result = json.loads(lines[0])
            assert "human_evaluation" in result
            assert result["human_evaluation"] is True


class TestGradeSingleResultEdgeCases:
    """Edge case tests for grade_single_result."""

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_with_empty_response(self, mock_get_eval_llm):
        """Test grading with empty model response."""
        from local_deep_research.benchmarks.graders import grade_single_result

        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = ""
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        result_data = {
            "problem": "Question",
            "correct_answer": "Answer",
            "response": "",
        }

        graded = grade_single_result(result_data)

        assert graded["is_correct"] is False

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_with_llm_no_invoke(self, mock_get_eval_llm):
        """Test grading when LLM doesn't have invoke method."""
        from local_deep_research.benchmarks.graders import grade_single_result

        # Create LLM without invoke method but with close for cleanup
        mock_llm = MagicMock(spec=["close"])
        mock_llm.__call__ = MagicMock(
            return_value="Extracted Answer: test\nCorrect: yes"
        )
        mock_get_eval_llm.return_value = mock_llm

        result_data = {
            "problem": "Q",
            "correct_answer": "A",
            "response": "R",
        }

        graded = grade_single_result(result_data)

        # Should still work via fallback
        assert "is_correct" in graded

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_with_chat_messages_attribute(self, mock_get_eval_llm):
        """Test grading with LLM that has chat_messages attribute."""
        from local_deep_research.benchmarks.graders import grade_single_result

        mock_llm = MagicMock()
        mock_llm.chat_messages = True  # Has this attribute
        mock_response = MagicMock()
        mock_response.content = "Extracted Answer: test\nCorrect: yes"
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        result_data = {
            "problem": "Q",
            "correct_answer": "A",
            "response": "R",
        }

        graded = grade_single_result(result_data)

        assert graded["is_correct"] is True
        # Should have called invoke with HumanMessage
        mock_llm.invoke.assert_called_once()

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_simpleqa_correct_no(self, mock_get_eval_llm):
        """Test SimpleQA grading with 'no' judgment."""
        from local_deep_research.benchmarks.graders import grade_single_result

        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = """
Extracted Answer: wrong answer
Reasoning: The model's answer is incorrect.
Correct: no
"""
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        result_data = {
            "problem": "What is 2+2?",
            "correct_answer": "4",
            "response": "The answer is 5.",
        }

        graded = grade_single_result(result_data, dataset_type="simpleqa")

        assert graded["is_correct"] is False

    @patch("local_deep_research.benchmarks.graders.get_evaluation_llm")
    def test_grade_preserves_settings_snapshot(self, mock_get_eval_llm):
        """Test that settings_snapshot is passed to get_evaluation_llm."""
        from local_deep_research.benchmarks.graders import grade_single_result

        mock_llm = MagicMock()
        mock_response = MagicMock()
        mock_response.content = "Extracted Answer: test\nCorrect: yes"
        mock_llm.invoke.return_value = mock_response
        mock_get_eval_llm.return_value = mock_llm

        settings_snapshot = {"llm.api_key": "test-key"}

        result_data = {
            "problem": "Q",
            "correct_answer": "A",
            "response": "R",
        }

        grade_single_result(result_data, settings_snapshot=settings_snapshot)

        # Verify settings_snapshot was passed
        mock_get_eval_llm.assert_called_once()
        call_args = mock_get_eval_llm.call_args
        assert (
            call_args[0][1] == settings_snapshot
            or call_args[1].get("settings_snapshot") == settings_snapshot
        )


class TestExtractAnswerEdgeCases:
    """Edge case tests for extract_answer_from_response."""

    def test_extract_handles_multiline_answer(self):
        """Test extraction of multiline answers."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = """Based on my research:

Exact Answer: This is a
multiline answer
Confidence: 90%
"""
        result = extract_answer_from_response(response, "browsecomp")

        # Should capture first line after "Exact Answer:"
        assert "This is a" in result["extracted_answer"]

    def test_extract_handles_special_characters(self):
        """Test extraction handles special characters."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = "The answer is: $100 (USD) [according to source]."
        result = extract_answer_from_response(response, "simpleqa")

        # Only numeric citations like [1], [2] are removed
        # Text-based brackets like [according to source] are preserved
        assert "$100" in result["extracted_answer"]

    def test_extract_empty_response(self):
        """Test extraction with empty response."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        result = extract_answer_from_response("", "simpleqa")

        assert result["extracted_answer"] == ""
        assert result["confidence"] == "100"

    def test_extract_browsecomp_no_exact_answer(self):
        """Test BrowseComp extraction without 'Exact Answer' marker."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = "The value is 42."
        result = extract_answer_from_response(response, "browsecomp")

        assert result["extracted_answer"] == "None"

    def test_extract_removes_multiple_citations(self):
        """Test that multiple citations are all removed."""
        from local_deep_research.benchmarks.graders import (
            extract_answer_from_response,
        )

        response = "First point [1], second point [2], third point [3][4][5]."
        result = extract_answer_from_response(response, "simpleqa")

        assert "[1]" not in result["extracted_answer"]
        assert "[2]" not in result["extracted_answer"]
        assert "[5]" not in result["extracted_answer"]
