"""
Test for citation formatting bug with multiple citations to the same source.

This test reproduces a bug where when a source is cited multiple times
and the citation format is not IEEE-style (e.g., using domain-based formats),
the formatting might incorrectly default to IEEE-style numbering.
"""

import pytest
from tests.test_utils import add_src_to_path

# Add src to path
add_src_to_path()

from local_deep_research.text_optimization.citation_formatter import (  # noqa: E402
    CitationFormatter,
    CitationMode,
)


class TestMultipleCitationsBug:
    """Test cases for the multiple citations formatting bug."""

    def test_domain_format_with_multiple_citations_same_source(self):
        """
        Test that domain-based citation format correctly handles multiple citations
        to the same source without reverting to IEEE-style numbering.
        """
        # Create test content with multiple citations to the same source
        content = """# Research Report

This is the first citation to arXiv [1]. Here is another citation
to the same arXiv paper [1]. And a third citation [1].

This is a citation to a different source [2].

## Sources

[1] Attention Is All You Need
URL: https://arxiv.org/abs/1706.03762

[2] BERT: Pre-training in Nature
URL: https://nature.com/articles/s41586-023-06785-z"""

        # Test with domain_hyperlinks mode
        formatter = CitationFormatter(mode=CitationMode.DOMAIN_HYPERLINKS)
        result = formatter.format_document(content)

        # All three citations to source 1 should be [arxiv.org]
        assert result.count("[arxiv.org]") == 3, (
            "Expected 3 instances of [arxiv.org] for multiple citations to same source, "
            f"but found {result.count('[arxiv.org]')}"
        )

        # One citation to source 2 should be [nature.com]
        assert result.count("[nature.com]") == 1, (
            "Expected 1 instance of [nature.com], "
            f"but found {result.count('[nature.com]')}"
        )

        # Should not contain any IEEE-style numbering [1] in body text
        body_text = result.split("## Sources")[0]
        assert "[1]" not in body_text, (
            "Found IEEE-style [1] citation in body when using domain format"
        )
        assert "[2]" not in body_text, (
            "Found IEEE-style [2] citation in body when using domain format"
        )

    def test_domain_id_format_with_multiple_citations_same_source(self):
        """
        Test domain_id format handles multiple citations to same source correctly.
        When there's only one source from a domain, it should use [domain.com],
        not [domain.com-1].
        """
        content = """# Research Report

First citation [1], second citation [1], third citation [1].
Citation to different paper from same domain [2].
Citation to different domain [3].

## Sources

[1] Paper One
URL: https://arxiv.org/abs/1706.03762

[2] Paper Two
URL: https://arxiv.org/abs/1810.04805

[3] Other Paper
URL: https://openai.com/blog/chatgpt"""

        formatter = CitationFormatter(mode=CitationMode.DOMAIN_ID_HYPERLINKS)
        result = formatter.format_document(content)

        # With smart numbering, since there are 2 arxiv papers, should use numbered format
        assert "[arxiv.org-1]" in result, (
            "Expected [arxiv.org-1] for first arxiv paper"
        )
        assert "[arxiv.org-2]" in result, (
            "Expected [arxiv.org-2] for second arxiv paper"
        )

        # Single source from openai.com should just be [openai.com]
        assert "[openai.com]" in result, (
            "Expected [openai.com] for single openai source"
        )
        assert "[openai.com-1]" not in result, (
            "Should not add -1 for single source from domain"
        )

        # Should not have any plain number citations in body
        body_text = result.split("## Sources")[0]
        assert "[1]" not in body_text, "Found IEEE-style [1] citation in body"
        assert "[2]" not in body_text, "Found IEEE-style [2] citation in body"
        assert "[3]" not in body_text, "Found IEEE-style [3] citation in body"

    def test_consecutive_citations_different_sources(self):
        """Test that consecutive citations to different sources are handled correctly."""
        content = """# Research Report

Here are three consecutive citations [1][2][3].

## Sources

[1] Paper One
URL: https://arxiv.org/abs/1111.1111

[2] Paper Two
URL: https://nature.com/articles/2222

[3] Paper Three
URL: https://science.org/doi/3333"""

        formatter = CitationFormatter(mode=CitationMode.DOMAIN_HYPERLINKS)
        result = formatter.format_document(content)

        # Should have three different domain citations
        assert "[arxiv.org]" in result
        assert "[nature.com]" in result
        assert "[science.org]" in result

        # Should not have any number citations in body
        body_text = result.split("## Sources")[0]
        assert "[1][2][3]" not in body_text
        assert "[1]" not in body_text
        assert "[2]" not in body_text
        assert "[3]" not in body_text

    def test_mixed_citation_formats_in_text(self):
        """Test handling of mixed citation formats like [1, 2, 3] vs [1][2][3]."""
        content = """# Research Report

Comma-separated citations [1, 2, 3].
Consecutive citations [1][2][3].
Single citations [1] and [2] and [3].

## Sources

[1] Paper One
URL: https://arxiv.org/abs/1111.1111

[2] Paper Two
URL: https://arxiv.org/abs/2222.2222

[3] Paper Three
URL: https://arxiv.org/abs/3333.3333"""

        formatter = CitationFormatter(
            mode=CitationMode.DOMAIN_ID_ALWAYS_HYPERLINKS
        )
        result = formatter.format_document(content)

        # All arxiv citations should be numbered consistently
        assert result.count("[arxiv.org-1]") == 3, (
            "Each citation to paper 1 should be [arxiv.org-1]"
        )
        assert result.count("[arxiv.org-2]") == 3, (
            "Each citation to paper 2 should be [arxiv.org-2]"
        )
        assert result.count("[arxiv.org-3]") == 3, (
            "Each citation to paper 3 should be [arxiv.org-3]"
        )

        # No plain numbers should remain in the body (before Sources section)
        body_text = result.split("## Sources")[0]
        for i in range(1, 4):
            assert f"[{i}]" not in body_text, (
                f"Found IEEE-style [{i}] citation in body"
            )

    def test_citation_format_preservation_in_sources_section(self):
        """Test that the Sources section format is preserved correctly."""
        content = """# Research Report

Citation one [1] and citation two [2].

## Sources

[1] Important Paper
URL: https://example.com/paper1

[2] Another Paper
URL: https://example.org/paper2

### References

Some additional references here."""

        formatter = CitationFormatter(mode=CitationMode.DOMAIN_HYPERLINKS)
        result = formatter.format_document(content)

        # Sources section should still have numbered format
        assert "[1] Important Paper" in result
        assert "[2] Another Paper" in result

        # But inline citations should be domain-based
        assert "[example.com]" in result
        assert "[example.org]" in result


if __name__ == "__main__":
    pytest.main([__file__, "-v"])