"""
Tests for report_generator.py

Tests cover:
- IntegratedReportGenerator initialization
- get_report_generator function
- Report structure determination
- Section generation
- Error handling
"""

from unittest.mock import Mock, patch


class TestGetReportGenerator:
    """Tests for get_report_generator function."""

    def test_get_report_generator_default(self):
        """Test get_report_generator returns IntegratedReportGenerator."""
        with patch(
            "local_deep_research.report_generator.IntegratedReportGenerator"
        ) as mock_class:
            mock_instance = Mock()
            mock_class.return_value = mock_instance

            from local_deep_research.report_generator import (
                get_report_generator,
            )

            result = get_report_generator()

            mock_class.assert_called_once_with(search_system=None)
            assert result == mock_instance

    def test_get_report_generator_with_search_system(self):
        """Test get_report_generator with search system."""
        mock_search_system = Mock()

        with patch(
            "local_deep_research.report_generator.IntegratedReportGenerator"
        ) as mock_class:
            mock_instance = Mock()
            mock_class.return_value = mock_instance

            from local_deep_research.report_generator import (
                get_report_generator,
            )

            get_report_generator(search_system=mock_search_system)

            mock_class.assert_called_once_with(search_system=mock_search_system)


class TestIntegratedReportGeneratorInit:
    """Tests for IntegratedReportGenerator initialization."""

    def test_init_with_search_system(self):
        """Test initialization with search system."""
        mock_search_system = Mock()
        mock_search_system.model = Mock()

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        assert generator.search_system == mock_search_system
        assert generator.model == mock_search_system.model
        assert generator.searches_per_section == 2

    def test_init_with_llm(self):
        """Test initialization with LLM only."""
        mock_llm = Mock()

        with patch(
            "local_deep_research.report_generator.AdvancedSearchSystem"
        ) as mock_search_class:
            mock_search = Mock()
            mock_search_class.return_value = mock_search

            from local_deep_research.report_generator import (
                IntegratedReportGenerator,
            )

            generator = IntegratedReportGenerator(llm=mock_llm)

            assert generator.model == mock_llm
            mock_search_class.assert_called_once_with(llm=mock_llm)

    def test_init_with_custom_searches_per_section(self):
        """Test initialization with custom searches per section."""
        mock_search_system = Mock()
        mock_search_system.model = Mock()

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(
            search_system=mock_search_system, searches_per_section=5
        )

        assert generator.searches_per_section == 5

    def test_init_search_system_overrides_llm(self):
        """Test search system's LLM is used when both provided."""
        mock_search_system = Mock()
        mock_search_system.model = Mock(name="search_system_model")
        mock_external_llm = Mock(name="external_llm")

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(
            search_system=mock_search_system, llm=mock_external_llm
        )

        # When both provided, external llm is used if specified
        assert generator.model == mock_external_llm


class TestDetermineReportStructure:
    """Tests for _determine_report_structure method."""

    def test_determine_structure_parses_response(self):
        """Test structure is parsed from LLM response."""
        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
   - Background | Historical context
2. Main Analysis
   - Data | Present findings
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 len(structure[0]["subsections"]) == 2
        assert structure[0]["subsections"][0]["name"] == "Overview"
        assert structure[0]["subsections"][0]["purpose"] == "Provide context"

    def test_determine_structure_removes_source_sections(self):
        """Test source-related sections are removed."""
        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. Sources and References
   - Citations | List 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"
        )

        # Should remove the Sources section
        assert len(structure) == 1
        assert structure[0]["name"] == "Introduction"

    def test_determine_structure_handles_subsection_without_purpose(self):
        """Test subsections without purpose get default."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm

        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Introduction
   - Overview
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 structure[0]["subsections"][0]["name"] == "Overview"
        assert "Overview" in structure[0]["subsections"][0]["purpose"]


class TestResearchAndGenerateSections:
    """Tests for _research_and_generate_sections method."""

    def test_research_sections_basic(self):
        """Test basic section research and generation."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.strategy.settings_snapshot = {"search.iterations": 3}
        mock_search_system.strategy.max_iterations = 3
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Section content here"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        structure = [
            {
                "name": "Introduction",
                "subsections": [
                    {"name": "Overview", "purpose": "Provide overview"}
                ],
            }
        ]
        initial_findings = {"questions_by_iteration": {}}

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

        assert "Introduction" in sections
        assert "Section content here" in sections["Introduction"]

    def test_research_sections_preserves_questions(self):
        """Test that questions from initial research are preserved."""
        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_search_system.questions_by_iteration = {}

        mock_strategy = Mock()
        mock_strategy.questions_by_iteration = {}
        mock_strategy.settings_snapshot = {"search.iterations": 3}
        mock_strategy.max_iterations = 3
        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": ["Q1", "Q2"]}}
        structure = [
            {
                "name": "Section",
                "subsections": [{"name": "Sub", "purpose": "Test"}],
            }
        ]

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

        # Questions should be set on strategy
        assert mock_strategy.questions_by_iteration == {"0": ["Q1", "Q2"]}

    def test_research_sections_creates_subsection_for_empty(self):
        """Test subsection is created if none provided."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.strategy.settings_snapshot = {"search.iterations": 3}
        mock_search_system.strategy.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!
            }
        ]
        initial_findings = {}

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

        # analyze_topic should still be called
        mock_search_system.analyze_topic.assert_called()


class TestFormatFinalReport:
    """Tests for _format_final_report method."""

    def test_format_final_report_structure(self):
        """Test final report has correct structure."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.all_links_of_system = []

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        sections = {"Introduction": "# Introduction\nContent here"}
        structure = [
            {
                "name": "Introduction",
                "subsections": [
                    {"name": "Overview", "purpose": "Provide overview"}
                ],
            }
        ]

        with patch(
            "local_deep_research.report_generator.importlib.import_module"
        ) as mock_import:
            mock_utilities = Mock()
            mock_utilities.search_utilities.format_links_to_markdown.return_value = "- [Link](url)"
            mock_import.return_value = mock_utilities

            result = generator._format_final_report(
                sections, structure, "test query"
            )

        assert "content" in result
        assert "metadata" in result
        assert "Table of Contents" in result["content"]
        assert "Introduction" in result["content"]
        assert "Sources" in result["content"]

    def test_format_final_report_metadata(self):
        """Test final report metadata."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.all_links_of_system = [
            {"link": "url1"},
            {"link": "url2"},
        ]

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        sections = {"Section1": "Content"}
        structure = [{"name": "Section1", "subsections": []}]

        with patch(
            "local_deep_research.report_generator.importlib.import_module"
        ) as mock_import:
            mock_utilities = Mock()
            mock_utilities.search_utilities.format_links_to_markdown.return_value = ""
            mock_import.return_value = mock_utilities

            result = generator._format_final_report(
                sections, structure, "test query"
            )

        assert result["metadata"]["initial_sources"] == 2
        assert result["metadata"]["sections_researched"] == 1
        assert result["metadata"]["query"] == "test query"
        assert "generated_at" in result["metadata"]


class TestGenerateReport:
    """Tests for generate_report method."""

    def test_generate_report_full_flow(self):
        """Test full report generation flow."""
        mock_search_system = Mock()
        mock_llm = Mock()
        mock_search_system.model = mock_llm
        mock_search_system.max_iterations = 3
        mock_search_system.strategy.settings_snapshot = {"search.iterations": 3}
        mock_search_system.strategy.max_iterations = 3
        mock_search_system.all_links_of_system = []

        # Mock LLM response for structure
        mock_response = Mock()
        mock_response.content = """
STRUCTURE
1. Introduction
   - Overview | Provide overview
END_STRUCTURE
"""
        mock_llm.invoke.return_value = mock_response

        # Mock search results
        mock_search_system.analyze_topic.return_value = {
            "current_knowledge": "Test content"
        }

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        initial_findings = {"current_knowledge": "Initial research"}

        with patch(
            "local_deep_research.report_generator.importlib.import_module"
        ) as mock_import:
            mock_utilities = Mock()
            mock_utilities.search_utilities.format_links_to_markdown.return_value = ""
            mock_import.return_value = mock_utilities

            result = generator.generate_report(initial_findings, "test query")

        assert "content" in result
        assert "metadata" in result


class TestGenerateErrorReport:
    """Tests for _generate_error_report method."""

    def test_generate_error_report(self):
        """Test error report generation."""
        mock_search_system = Mock()
        mock_search_system.model = Mock()

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        result = generator._generate_error_report(
            "test query", "Something went wrong"
        )

        assert "ERROR REPORT" in result
        assert "test query" in result
        assert "Something went wrong" in result


class TestGenerateSectionsDeprecated:
    """Tests for deprecated _generate_sections method."""

    def test_generate_sections_returns_empty(self):
        """Test deprecated _generate_sections returns empty dict."""
        mock_search_system = Mock()
        mock_search_system.model = Mock()

        from local_deep_research.report_generator import (
            IntegratedReportGenerator,
        )

        generator = IntegratedReportGenerator(search_system=mock_search_system)

        result = generator._generate_sections({}, {}, [], "query")

        assert result == {}
