"""Tests for citation post-processing, focused on strip_llm_sources_section.

Covers:
- Heading-based stripping (## Sources, ## References) — regression
- Inline italic source lists (*Sources: [N] ...*)
- Plain inline source lists (Sources: [N] ...)
- Horizontal rule preceding inline sources (--- + Sources:)
- Combined heading + inline sources
- No sources at all (no-op)
"""

from __future__ import annotations

from foundry_mcp.core.research.workflows.deep_research.phases._citation_postprocess import (
    strip_llm_sources_section,
)


# ---------------------------------------------------------------------------
# Heading-based stripping (existing behavior — regression tests)
# ---------------------------------------------------------------------------


class TestStripHeadingBasedSources:
    """Verify existing ## Sources / ## References stripping still works."""

    def test_strip_sources_heading(self):
        report = (
            "# Report\n\n"
            "Some content here.\n\n"
            "## Sources\n\n"
            "[1] [Title](https://example.com)\n"
            "[2] [Title2](https://example2.com)\n"
        )
        result = strip_llm_sources_section(report)
        assert "## Sources" not in result
        assert "[1]" not in result
        assert "Some content here." in result

    def test_strip_references_heading(self):
        report = (
            "# Report\n\n"
            "Content.\n\n"
            "## References\n\n"
            "[1] Author (2025). Title.\n"
        )
        result = strip_llm_sources_section(report)
        assert "## References" not in result
        assert "Content." in result

    def test_preserves_content_after_next_heading(self):
        report = (
            "# Report\n\n"
            "Content.\n\n"
            "## Sources\n\n"
            "[1] Source\n\n"
            "## Appendix\n\n"
            "Extra content.\n"
        )
        result = strip_llm_sources_section(report)
        assert "## Sources" not in result
        assert "## Appendix" in result
        assert "Extra content." in result


# ---------------------------------------------------------------------------
# Inline source list stripping (new behavior)
# ---------------------------------------------------------------------------


class TestStripInlineSources:
    """Strip inline *Sources: [N] ...* paragraphs generated by the LLM."""

    def test_strip_italic_inline_sources(self):
        report = (
            "# Report\n\n"
            "Some analysis here.\n\n"
            "*Sources: [2] [LendingTree](https://example.com), "
            "[3] [NerdWallet](https://example2.com)*\n"
        )
        result = strip_llm_sources_section(report)
        assert "*Sources:" not in result
        assert "LendingTree" not in result
        assert "Some analysis here." in result

    def test_strip_plain_inline_sources(self):
        report = (
            "# Report\n\n"
            "Body text.\n\n"
            "Sources: [1] First source, [2] Second source\n"
        )
        result = strip_llm_sources_section(report)
        assert "Sources: [1]" not in result
        assert "Body text." in result

    def test_strip_inline_sources_with_hr_rule(self):
        report = (
            "# Report\n\n"
            "Content paragraph.\n\n"
            "---\n"
            "*Sources: [1] [Title](https://example.com), [2] [Title2](https://example2.com)*\n"
        )
        result = strip_llm_sources_section(report)
        assert "---" not in result
        assert "*Sources:" not in result
        assert "Content paragraph." in result

    def test_strip_inline_sources_preserves_following_content(self):
        report = (
            "# Report\n\n"
            "Main content.\n\n"
            "*Sources: [1] [Source](https://example.com)*\n\n"
            "## Footnotes\n\n"
            "Some footnotes.\n"
        )
        result = strip_llm_sources_section(report)
        assert "*Sources:" not in result
        assert "## Footnotes" in result
        assert "Some footnotes." in result

    def test_strip_multiline_inline_sources(self):
        """Inline sources spanning multiple lines before next blank line."""
        report = (
            "# Report\n\n"
            "Analysis.\n\n"
            "*Sources: [1] [First](https://a.com),\n"
            "[2] [Second](https://b.com),\n"
            "[3] [Third](https://c.com)*\n\n"
            "## Next Section\n\n"
            "More content.\n"
        )
        result = strip_llm_sources_section(report)
        assert "*Sources:" not in result
        assert "[First]" not in result
        assert "## Next Section" in result
        assert "More content." in result


# ---------------------------------------------------------------------------
# Combined heading + inline sources
# ---------------------------------------------------------------------------


class TestStripBothHeadingAndInline:
    """Reports with both ## Sources heading AND inline sources paragraph."""

    def test_strip_both_heading_and_inline(self):
        report = (
            "# Report\n\n"
            "Content.\n\n"
            "*Sources: [1] [Inline](https://inline.com)*\n\n"
            "## Sources\n\n"
            "[1] [Heading Source](https://heading.com)\n"
        )
        result = strip_llm_sources_section(report)
        assert "## Sources" not in result
        assert "*Sources:" not in result
        assert "[Inline]" not in result
        assert "[Heading Source]" not in result
        assert "Content." in result


# ---------------------------------------------------------------------------
# No sources (no-op)
# ---------------------------------------------------------------------------


class TestNoSources:
    """Report with no source sections at all — should be returned unchanged."""

    def test_no_sources_noop(self):
        report = "# Report\n\nJust some content.\n\n## Conclusion\n\nDone.\n"
        result = strip_llm_sources_section(report)
        assert result == report

    def test_empty_report(self):
        assert strip_llm_sources_section("") == ""

    def test_report_with_sources_word_in_body(self):
        """'Sources' in body text should not be stripped."""
        report = (
            "# Report\n\n"
            "We consulted multiple sources for this analysis.\n\n"
            "## Conclusion\n\n"
            "Done.\n"
        )
        result = strip_llm_sources_section(report)
        assert "We consulted multiple sources" in result
