"""Tests for citation post-processing.

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)
- Citation year-reference filtering (_CITATION_RE skips [1900]-[2099])
"""

from __future__ import annotations

from foundry_mcp.core.research.models.deep_research import DeepResearchState
from foundry_mcp.core.research.workflows.deep_research.phases._citation_postprocess import (
    _CITATION_RE,
    cleanup_citations,
    extract_cited_numbers,
    renumber_citations,
    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


# ---------------------------------------------------------------------------
# cleanup_citations: strip-only behavior (no renumber, no bibliography)
# ---------------------------------------------------------------------------


class TestCleanupCitationsStripOnly:
    """Verify cleanup_citations strips LLM sources and dangling but does NOT
    renumber or append bibliography."""

    @staticmethod
    def _make_state(n: int = 3) -> DeepResearchState:
        state = DeepResearchState(original_query="test")
        for i in range(n):
            state.add_source(title=f"Source {i+1}", url=f"https://s{i+1}.example.com")
        return state

    def test_strips_heading_sources_section(self):
        state = self._make_state()
        report = (
            "# Report\n\n"
            "Finding [1] and [2].\n\n"
            "## Sources\n\n"
            "[1] [Source 1](https://s1.example.com)\n"
            "[2] [Source 2](https://s2.example.com)\n"
        )
        cleaned, _ = cleanup_citations(report, state)
        assert "## Sources" not in cleaned
        assert "Finding [1] and [2]." in cleaned

    def test_strips_inline_sources(self):
        state = self._make_state()
        report = (
            "# Report\n\n"
            "Analysis [1].\n\n"
            "*Sources: [1] [Source 1](https://s1.example.com)*\n"
        )
        cleaned, _ = cleanup_citations(report, state)
        assert "*Sources:" not in cleaned
        assert "Analysis [1]." in cleaned

    def test_removes_dangling_keeps_valid(self):
        state = self._make_state(2)
        report = "Finding [1], [2], and [99]."
        cleaned, meta = cleanup_citations(report, state)
        assert "[1]" in cleaned
        assert "[2]" in cleaned
        assert "[99]" not in cleaned
        assert meta["dangling_citations_removed"] == 1

    def test_no_renumber_out_of_order(self):
        """Out-of-order citations remain out of order after cleanup."""
        state = self._make_state(3)
        report = "First [3] then [1]."
        cleaned, _ = cleanup_citations(report, state)
        assert "First [3] then [1]." in cleaned

    def test_no_bibliography_appended(self):
        state = self._make_state(2)
        report = "Finding [1]."
        cleaned, _ = cleanup_citations(report, state)
        assert "## Sources" not in cleaned
        assert "## References" not in cleaned

    def test_combined_strip_and_dangling(self):
        """Both LLM sources section AND dangling citations are handled."""
        state = self._make_state(2)
        report = (
            "# Report\n\n"
            "Finding [1] and [50] are important.\n\n"
            "## Sources\n\n"
            "- LLM generated sources\n"
        )
        cleaned, meta = cleanup_citations(report, state)
        assert "## Sources" not in cleaned
        assert "LLM generated" not in cleaned
        assert "[1]" in cleaned
        assert "[50]" not in cleaned
        assert meta["dangling_citations_removed"] == 1


# ---------------------------------------------------------------------------
# Citation year-reference filtering (_CITATION_RE)
# ---------------------------------------------------------------------------


class TestCitationYearFiltering:
    """Verify _CITATION_RE skips year references [1900]–[2099]."""

    def test_matches_low_citation_numbers(self):
        """[1], [2], [99] should be matched as citations."""
        for text, expected in [("[1]", "1"), ("[2]", "2"), ("[99]", "99")]:
            matches = _CITATION_RE.findall(text)
            assert matches == [expected], f"{text} should match"

    def test_matches_high_citation_numbers(self):
        """[100], [500] should be matched (valid high citation numbers)."""
        for text, expected in [("[100]", "100"), ("[500]", "500")]:
            matches = _CITATION_RE.findall(text)
            assert matches == [expected], f"{text} should match"

    def test_skips_year_references(self):
        """[2026], [2025], [1999] should NOT be matched."""
        for text in ["[2026]", "[2025]", "[1999]", "[1900]", "[2099]"]:
            matches = _CITATION_RE.findall(text)
            assert matches == [], f"{text} should NOT match"

    def test_matches_outside_year_range(self):
        """[1899], [2100] are outside the year range — should match."""
        for text, expected in [("[1899]", "1899"), ("[2100]", "2100")]:
            matches = _CITATION_RE.findall(text)
            assert matches == [expected], f"{text} should match"

    def test_year_not_extracted_by_extract_cited_numbers(self):
        """Year references in body text should not be extracted."""
        report = "Published in [2026] by researchers, citing [1] and [2]."
        numbers = extract_cited_numbers(report)
        assert 2026 not in numbers
        assert numbers == {1, 2}

    def test_renumber_preserves_year_references(self):
        """renumber_citations should leave [2026] untouched."""
        state = DeepResearchState(original_query="test")
        state.add_source(title="Source A", url="https://a.com")
        state.add_source(title="Source B", url="https://b.com")

        report = "Finding [2] then [1] in [2026]."
        result, remap = renumber_citations(report, state)
        assert "[2026]" in result
        assert "[1]" in result  # renumbered from [2]
        assert "[2]" in result  # renumbered from [1]
        assert 2026 not in remap
