"""
Extended tests for report_generator.py

Tests cover edge cases and scenarios not covered in the base test file:
- Structure parsing edge cases
- Source section removal with various keywords
- Malformed LLM response handling
- Subsection parsing edge cases
- Max iterations modification and restoration
- Question preservation across sections
"""

from unittest.mock import Mock


class TestDetermineReportStructureMarkers:
    """Tests for structure marker parsing in _determine_report_structure."""

    def test_parses_structure_without_end_marker(self):
        """Test parsing when END_STRUCTURE is missing."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Introduction
   - Overview | Provide context
2. Analysis
   - Details | Explain findings
"""  # No END_STRUCTURE
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert len(structure) == 2
        assert structure[0]["name"] == "Introduction"
        assert structure[1]["name"] == "Analysis"

    def test_parses_structure_without_start_marker(self):
        """Test parsing when STRUCTURE marker is missing."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
1. Introduction
   - Overview | Provide context
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        # Should still parse sections
        assert len(structure) >= 1

    def test_handles_numbered_sections_various_digits(self):
        """Test parsing sections with various digit numbers."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. First
2. Second
3. Third
9. Ninth
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert len(structure) == 4
        assert structure[3]["name"] == "Ninth"

    def test_ignores_lines_without_section_format(self):
        """Test that non-section lines are ignored."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
Here is the report structure:
1. Introduction
   - Overview | Context
Some random text here
2. Conclusion
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert len(structure) == 2
        assert structure[0]["name"] == "Introduction"
        assert structure[1]["name"] == "Conclusion"


class TestRemoveSourceSectionsKeywords:
    """Tests for source section removal with various keywords."""

    def test_removes_citation_section(self):
        """Test removes section with 'citation' keyword."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Introduction
   - Overview | Context
2. Citations and References
   - Bibliography | List all
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert len(structure) == 1
        assert structure[0]["name"] == "Introduction"

    def test_removes_bibliography_section(self):
        """Test removes section with 'bibliography' keyword."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Main Content
   - Details | Explain
2. Bibliography
   - Works Cited | References
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert len(structure) == 1
        assert structure[0]["name"] == "Main Content"

    def test_removes_reference_section(self):
        """Test removes section with 'reference' keyword."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Analysis
   - Data | Present findings
2. References
   - Links | All sources
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert len(structure) == 1
        assert structure[0]["name"] == "Analysis"

    def test_only_removes_last_source_section(self):
        """Test only last section is checked for source keywords."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Source Code Analysis
   - Details | Analyze source code
2. Conclusion
   - Summary | Wrap up
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        # "Source Code Analysis" should NOT be removed (not last section)
        assert len(structure) == 2
        assert structure[0]["name"] == "Source Code Analysis"

    def test_case_insensitive_source_detection(self):
        """Test source keyword detection is case insensitive."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Analysis
   - Data | Present findings
2. SOURCES AND CITATIONS
   - Links | All sources
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert len(structure) == 1


class TestHandleMalformedResponse:
    """Tests for handling malformed LLM responses."""

    def test_handles_empty_response(self):
        """Test handles empty LLM response."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

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

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert structure == []

    def test_handles_whitespace_only_response(self):
        """Test handles whitespace-only LLM response."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = "   \n\n   \t\t   "
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert structure == []

    def test_handles_response_with_only_markers(self):
        """Test handles response with only STRUCTURE markers."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = "STRUCTURE\nEND_STRUCTURE"
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert structure == []

    def test_handles_subsection_before_section(self):
        """Test handles subsection appearing before any section."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
   - Orphan Subsection | No parent
1. First Section
   - Valid Subsection | Has parent
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        # Orphan subsection should be ignored
        assert len(structure) == 1
        assert len(structure[0]["subsections"]) == 1


class TestSubsectionParsing:
    """Tests for subsection parsing edge cases."""

    def test_subsection_with_multiple_pipes(self):
        """Test subsection with multiple pipe characters."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Section
   - Name with | pipe | characters | purpose here
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        # Only first pipe should be used as separator
        assert len(structure[0]["subsections"]) == 1
        assert structure[0]["subsections"][0]["name"] == "Name with"

    def test_subsection_with_empty_name(self):
        """Test subsection with empty name is ignored."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Section
   - | purpose only
   - Valid Name | purpose
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        # Empty name subsection should have empty string as name
        # but the parsing should still work
        assert len(structure[0]["subsections"]) >= 1

    def test_many_subsections(self):
        """Test section with many subsections."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        subsections = "\n".join(
            [f"   - Subsection {i} | Purpose {i}" for i in range(20)]
        )
        mock_response = Mock()
        mock_response.content = f"""
STRUCTURE
1. Large Section
{subsections}
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        findings = {"current_knowledge": "Test content " * 100}
        structure = generator._determine_report_structure(
            findings, "test query"
        )

        assert len(structure[0]["subsections"]) == 20


class TestMaxIterationsModificationAndRestore:
    """Tests for max_iterations modification during section research."""

    def test_max_iterations_set_to_one_during_search(self):
        """Test max_iterations is set to 1 during subsection search."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 5
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Sub", "purpose": "Test"}],
            }
        ]
        initial_findings = {}

        # Capture max_iterations during analyze_topic call
        captured_max_iterations = []

        def capture_max(*args, **kwargs):
            captured_max_iterations.append(mock_search_system.max_iterations)
            return {"current_knowledge": "Content"}

        mock_search_system.analyze_topic.side_effect = capture_max

        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        # During the call, max_iterations should have been 1
        assert 1 in captured_max_iterations

    def test_max_iterations_restored_after_search(self):
        """Test max_iterations is restored after section research."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 7
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Sub", "purpose": "Test"}],
            }
        ]
        initial_findings = {}

        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        # Should be restored to original value
        assert mock_search_system.max_iterations == 7

    def test_max_iterations_restored_even_with_multiple_sections(self):
        """Test max_iterations is restored after multiple sections."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 10
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Section 1",
                "subsections": [
                    {"name": "Sub 1", "purpose": "Test 1"},
                    {"name": "Sub 2", "purpose": "Test 2"},
                ],
            },
            {
                "name": "Section 2",
                "subsections": [{"name": "Sub 3", "purpose": "Test 3"}],
            },
        ]
        initial_findings = {}

        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        assert mock_search_system.max_iterations == 10


class TestPreserveQuestionsFromInitial:
    """Tests for preserving questions from initial research."""

    def test_questions_set_on_search_system(self):
        """Test questions are set on search system."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.questions_by_iteration = {}
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        initial_findings = {
            "questions_by_iteration": {
                1: ["Q1", "Q2"],
                2: ["Q3"],
            }
        }
        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Sub", "purpose": "Test"}],
            }
        ]

        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        assert mock_search_system.questions_by_iteration == {
            1: ["Q1", "Q2"],
            2: ["Q3"],
        }

    def test_questions_set_on_strategy(self):
        """Test questions are set on strategy."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        mock_strategy = Mock()
        mock_strategy.questions_by_iteration = {}
        mock_search_system.strategy = mock_strategy

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        initial_findings = {"questions_by_iteration": {0: ["Initial Q"]}}
        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Sub", "purpose": "Test"}],
            }
        ]

        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        assert mock_strategy.questions_by_iteration == {0: ["Initial Q"]}

    def test_handles_empty_questions(self):
        """Test handles empty questions gracefully."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        initial_findings = {"questions_by_iteration": {}}
        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Sub", "purpose": "Test"}],
            }
        ]

        # Should not raise
        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

    def test_handles_missing_questions_key(self):
        """Test handles missing questions_by_iteration key."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        initial_findings = {}  # No questions_by_iteration key
        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Sub", "purpose": "Test"}],
            }
        ]

        # Should not raise
        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )


class TestAutoGenerateSubsections:
    """Tests for auto-generating subsections when none provided."""

    def test_creates_subsection_from_section_name(self):
        """Test subsection is created from section name when empty."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Introduction",
                "subsections": [],  # Empty - should auto-generate
            }
        ]
        initial_findings = {}

        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        # analyze_topic should be called with section-level query
        assert mock_search_system.analyze_topic.called
        call_args = mock_search_system.analyze_topic.call_args
        assert "Introduction" in call_args[0][0]

    def test_section_name_with_pipe_creates_subsection(self):
        """Test section name with pipe is parsed into subsection."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Overview | Provide general context",
                "subsections": [],
            }
        ]
        initial_findings = {}

        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        # Should have been parsed and search called
        assert mock_search_system.analyze_topic.called

    def test_handles_limited_knowledge_result(self):
        """Test handles when analyze_topic returns no current_knowledge."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": None
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Sub", "purpose": "Test"}],
            }
        ]
        initial_findings = {}

        sections = generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        # Should contain fallback message
        assert "Limited information" in sections["Section"]


class TestResearchAndGenerateSectionsEdgeCases:
    """Additional edge case tests for _research_and_generate_sections."""

    def test_multiple_subsections_adds_headers(self):
        """Test multiple subsections get headers."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Main Section",
                "subsections": [
                    {"name": "First Sub", "purpose": "First purpose"},
                    {"name": "Second Sub", "purpose": "Second purpose"},
                ],
            }
        ]
        initial_findings = {}

        sections = generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        # Both subsections should have headers
        assert "## First Sub" in sections["Main Section"]
        assert "## Second Sub" in sections["Main Section"]

    def test_single_subsection_no_extra_header(self):
        """Test single subsection doesn't get extra header."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Only Sub", "purpose": "Purpose"}],
            }
        ]
        initial_findings = {}

        sections = generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        # Single subsection shouldn't have ## header (section header is # Section)
        assert "## Only Sub" not in sections["Section"]

    def test_context_includes_other_sections(self):
        """Test query includes context about other sections."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3

        captured_queries = []

        def capture_query(query):
            captured_queries.append(query)
            return {"current_knowledge": "Content"}

        mock_search_system.analyze_topic.side_effect = capture_query

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Section A",
                "subsections": [{"name": "Sub A", "purpose": "Purpose A"}],
            },
            {
                "name": "Section B",
                "subsections": [{"name": "Sub B", "purpose": "Purpose B"}],
            },
        ]
        initial_findings = {}

        generator._research_and_generate_sections(
            initial_findings, structure, "test query"
        )

        # Query for Section A should mention Section B as other section
        assert "Section B" in captured_queries[0]
        # Query for Section B should mention Section A as other section
        assert "Section A" in captured_queries[1]
