"""Tests for citation formatter functionality."""

import pytest
from local_deep_research.text_optimization import (
    CitationFormatter,
    CitationMode,
    LaTeXExporter,
    QuartoExporter,
    RISExporter,
)


class TestCitationFormatter:
    """Test cases for CitationFormatter class."""

    @pytest.fixture
    def sample_content(self):
        """Sample markdown content with citations."""
        return """# Research Report

This is a research report about artificial intelligence [1]. Recent advances in
machine learning [2] have shown significant progress. The work by researchers
at DeepMind [3] has been particularly influential.

## Key Findings

The analysis reveals several important points:
- Deep learning models are becoming more efficient [1]
- Transfer learning enables better generalization [2]
- Large language models show emergent capabilities [3]
- Multiple studies confirm these findings [1][2][3]
- Some research combines different approaches [2][3]
- Comprehensive surveys cover all aspects [1, 2, 3]
- Recent work builds on earlier findings [2, 3]

## Sources

[1] Understanding Deep Learning
    URL: https://arxiv.org/abs/2104.12345

[2] Transfer Learning: A Survey
    URL: https://www.nature.com/articles/s41586-021-03819-2

[3] Emergent Abilities of Large Language Models
    URL: https://openai.com/research/emergent-abilities
"""

    @pytest.fixture
    def content_without_urls(self):
        """Sample content with citations but no URLs."""
        return """# Report

Some findings [1] and more [2].

## Sources

[1] First Source
[2] Second Source
"""

    def test_number_hyperlinks_mode(self, sample_content):
        """Test citation formatting with number hyperlinks."""
        formatter = CitationFormatter(CitationMode.NUMBER_HYPERLINKS)
        result = formatter.format_document(sample_content)

        # Check single citations are formatted correctly
        assert "[[1]](https://arxiv.org/abs/2104.12345)" in result
        assert (
            "[[2]](https://www.nature.com/articles/s41586-021-03819-2)"
            in result
        )
        assert "[[3]](https://openai.com/research/emergent-abilities)" in result

        # Check multiple consecutive citations
        assert (
            "[[1]](https://arxiv.org/abs/2104.12345)[[2]](https://www.nature.com/articles/s41586-021-03819-2)[[3]](https://openai.com/research/emergent-abilities)"
            in result
        )

        # Check comma-separated citations
        assert (
            "[[1]](https://arxiv.org/abs/2104.12345)[[2]](https://www.nature.com/articles/s41586-021-03819-2)[[3]](https://openai.com/research/emergent-abilities)"
            in result
        )

        # Ensure sources section is unchanged
        assert "## Sources" in result
        assert "[1] Understanding Deep Learning" in result

    def test_domain_hyperlinks_mode(self, sample_content):
        """Test citation formatting with domain hyperlinks."""
        formatter = CitationFormatter(CitationMode.DOMAIN_HYPERLINKS)
        result = formatter.format_document(sample_content)

        # Check domain formatting
        assert "[[arxiv.org]](https://arxiv.org/abs/2104.12345)" in result
        assert (
            "[[nature.com]](https://www.nature.com/articles/s41586-021-03819-2)"
            in result
        )
        assert (
            "[[openai.com]](https://openai.com/research/emergent-abilities)"
            in result
        )

        # Check multiple citations work
        assert (
            "[[arxiv.org]](https://arxiv.org/abs/2104.12345)[[nature.com]](https://www.nature.com/articles/s41586-021-03819-2)[[openai.com]](https://openai.com/research/emergent-abilities)"
            in result
        )

    def test_no_hyperlinks_mode(self, sample_content):
        """Test citation formatting with no hyperlinks."""
        formatter = CitationFormatter(CitationMode.NO_HYPERLINKS)
        result = formatter.format_document(sample_content)

        # Content should be unchanged
        assert result == sample_content
        assert "[[1]]" not in result
        assert "](https://" not in result

    def test_citations_without_urls(self, content_without_urls):
        """Test handling of citations without URLs."""
        formatter = CitationFormatter(CitationMode.NUMBER_HYPERLINKS)
        result = formatter.format_document(content_without_urls)

        # Citations without URLs should remain unchanged
        assert "[1]" in result
        assert "[2]" in result
        assert "[[1]]" not in result
        assert "[[2]]" not in result

    def test_no_sources_section(self):
        """Test handling of content without sources section."""
        content = "This is text with [1] citation but no sources."
        formatter = CitationFormatter(CitationMode.NUMBER_HYPERLINKS)
        result = formatter.format_document(content)

        # Should return unchanged content
        assert result == content

    def test_mixed_citation_formats(self):
        """Test handling of mixed citation formats."""
        content = """Text with various formats:
- Single: [1]
- Multiple: [1][2]
- Comma-separated: [1, 2, 3]
- With spaces: [1,2,3] and [1, 2]

## Sources

[1] Source One
    URL: https://example1.com
[2] Source Two
    URL: https://example2.com
[3] Source Three
    URL: https://example3.com
"""
        formatter = CitationFormatter(CitationMode.NUMBER_HYPERLINKS)
        result = formatter.format_document(content)

        # All formats should be converted
        assert "[[1]](https://example1.com)" in result
        assert (
            "[[1]](https://example1.com)[[2]](https://example2.com)" in result
        )
        assert (
            "[[1]](https://example1.com)[[2]](https://example2.com)[[3]](https://example3.com)"
            in result
        )

    def test_domain_extraction(self):
        """Test domain extraction from various URLs."""
        formatter = CitationFormatter(CitationMode.DOMAIN_HYPERLINKS)

        # Test various URL formats
        test_cases = [
            ("https://arxiv.org/abs/1234", "arxiv.org"),
            ("https://www.nature.com/articles/123", "nature.com"),
            ("https://github.com/user/repo", "github.com"),
            ("https://subdomain.example.com/path", "example.com"),
            ("https://example.co.uk/path", "co.uk"),
        ]

        for url, expected_domain in test_cases:
            domain = formatter._extract_domain(url)
            assert domain == expected_domain

    def test_edge_cases(self):
        """Test edge cases and malformed content."""
        formatter = CitationFormatter(CitationMode.NUMBER_HYPERLINKS)

        # Empty content
        assert formatter.format_document("") == ""

        # Content with only sources section
        content = """## Sources
[1] Test
    URL: https://test.com
"""
        result = formatter.format_document(content)
        assert result == content

        # Malformed citations
        content = """Text with [a] non-numeric [citation].

## Sources
[a] Invalid citation
"""
        result = formatter.format_document(content)
        assert "[a]" in result  # Should remain unchanged
        assert "[citation]" in result

    def test_source_word_pattern_detection(self):
        """Test detection and conversion of 'Source X' patterns."""
        content = """# Research Report

As mentioned in Source 1, the data shows improvements.
According to source 2, this is confirmed.
The Source 3 provides additional evidence.

## Sources

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

[2] Second Study
    URL: https://example.com/paper2

[3] Third Analysis
    URL: https://example.com/paper3
"""
        formatter = CitationFormatter(CitationMode.NUMBER_HYPERLINKS)
        result = formatter.format_document(content)

        # Check that "Source X" patterns are converted to hyperlinks
        assert "[[1]](https://example.com/paper1)" in result
        assert "[[2]](https://example.com/paper2)" in result
        assert "[[3]](https://example.com/paper3)" in result

        # Original "Source X" text should be replaced
        assert "Source 1" not in result or "[[1]]" in result
        assert "source 2" not in result or "[[2]]" in result

    def test_source_pattern_edge_cases(self):
        """Test that Source pattern converts to broken references when no matching source."""
        content = """# Financial Report

Source 12 of income comes from investments.
The power source 5 failed during testing.
Open Source 100 projects were analyzed.
Valid reference from Source 1 here.

## Sources

[1] Financial Analysis
    URL: https://example.com/finance
"""
        formatter = CitationFormatter(CitationMode.NUMBER_HYPERLINKS)
        result = formatter.format_document(content)

        # These get converted to broken reference markers [X] since no matching sources
        assert "[12] of income" in result  # Converted but not hyperlinked
        assert "[5] failed" in result  # Converted but not hyperlinked
        assert "[100] projects" in result  # Converted but not hyperlinked

        # Valid Source 1 should be hyperlinked
        assert "[[1]](https://example.com/finance) here" in result

        # The broken references should NOT be hyperlinked
        assert "[[12]]" not in result
        assert "[[5]]" not in result
        assert "[[100]]" not in result

    def test_mixed_citation_formats_with_source_pattern(self):
        """Test document with both [X] and Source X patterns."""
        content = """# Research Report

Traditional citation [1] appears here.
Then Source 2 is mentioned directly.
Mixed with [3] and source 1 again.

## Sources

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

[2] Second Paper
    URL: https://example.com/paper2

[3] Third Paper
    URL: https://example.com/paper3
"""
        formatter = CitationFormatter(CitationMode.NUMBER_HYPERLINKS)
        result = formatter.format_document(content)

        # All formats should be converted consistently
        assert (
            result.count("[[1]](https://example.com/paper1)") >= 2
        )  # [1] and source 1
        assert "[[2]](https://example.com/paper2)" in result  # Source 2
        assert "[[3]](https://example.com/paper3)" in result  # [3]

    def test_domain_id_hyperlinks_mode(self, sample_content):
        """Test citation formatting with domain IDs for repeated domains."""
        # Add more sources from same domains to test ID functionality
        content_with_repeated_domains = """# Research Report

Citations from GitHub [1] and more from GitHub [2].
ArXiv papers [3] and another ArXiv paper [4].
A Nature article [5] and OpenAI research [6].

## Sources

[1] First GitHub Project
    URL: https://github.com/user/project1

[2] Second GitHub Project
    URL: https://github.com/user/project2

[3] First ArXiv Paper
    URL: https://arxiv.org/abs/2024.1111

[4] Second ArXiv Paper
    URL: https://arxiv.org/abs/2024.2222

[5] Nature Article
    URL: https://www.nature.com/articles/s41586-024-1234

[6] OpenAI Research
    URL: https://openai.com/research/gpt4
"""
        formatter = CitationFormatter(CitationMode.DOMAIN_ID_HYPERLINKS)
        result = formatter.format_document(content_with_repeated_domains)

        # Check that repeated domains get hyphen-separated IDs
        assert "[[github.com-1]](https://github.com/user/project1)" in result
        assert "[[github.com-2]](https://github.com/user/project2)" in result
        assert "[[arxiv.org-1]](https://arxiv.org/abs/2024.1111)" in result
        assert "[[arxiv.org-2]](https://arxiv.org/abs/2024.2222)" in result

        # Check that single domain citations don't get IDs
        assert (
            "[[nature.com]](https://www.nature.com/articles/s41586-024-1234)"
            in result
        )
        assert "[[openai.com]](https://openai.com/research/gpt4)" in result

    def test_real_world_example(self):
        """Test with real-world example from Local Deep Research output."""
        content = """# Quick Research Summary

Query: local deep research

Local Deep Research represents a new generation of AI-powered research assistants [1][3][11]. It utilizes any LLM hosted by either Ollama or LMStudio [2]. The concept was introduced as DeepSearcher [4] and further explained on Medium [7]. Web searches are conducted using DuckDuckGo [5], and the system incorporates tools like Milvus [6]. Updates were discussed on Reddit [10] and demonstrations on YouTube [8][12]. The Ollama variant [9] provides an alternative approach.

## Sources

[1] LearningCircuit/local-deep-research- GitHub
    URL: https://github.com/LearningCircuit/local-deep-research

[2] langchain-ai/local-deep-researcher - GitHub
    URL: https://github.com/langchain-ai/local-deep-researcher

[3] local-deep-research- PyPI
    URL: https://pypi.org/project/local-deep-research/

[4] Introducing DeepSearcher
    URL: https://zilliz.com/blog/introduce-deepsearcher

[5] A local LLM research assistant
    URL: https://www.reddit.com/r/LocalLLaMA/comments/1ilkosp/

[6] Introducing DeepSearcher - Milvus
    URL: https://milvus.io/blog/introduce-deepsearcher.md

[7] DeepSearcher on Medium
    URL: https://milvusio.medium.com/deepsearcher

[8] Private AI Deep Research - YouTube
    URL: https://m.youtube.com/watch?v=4tDqV__jzKY

[9] Ollama Deep Research
    URL: https://apidog.com/blog/ollama-deep-research/

[10] Local Deep Research Update
    URL: https://www.reddit.com/r/LocalLLaMA/comments/1j79obx/

[11] local-deep-research 0.3.3 - PyPI
    URL: https://pypi.org/project/local-deep-research/0.3.3/

[12] Building a local deep researcher - YouTube
    URL: https://www.youtube.com/watch?v=sGUjmyfof4Q
"""
        formatter = CitationFormatter(CitationMode.DOMAIN_ID_HYPERLINKS)
        result = formatter.format_document(content)

        # Check multiple citations from same domain get IDs
        assert (
            "[[github.com-1]](https://github.com/LearningCircuit/local-deep-research)"
            in result
        )
        assert (
            "[[github.com-2]](https://github.com/langchain-ai/local-deep-researcher)"
            in result
        )
        assert (
            "[[pypi.org-1]](https://pypi.org/project/local-deep-research/)"
            in result
        )
        assert (
            "[[pypi.org-2]](https://pypi.org/project/local-deep-research/0.3.3/)"
            in result
        )
        assert (
            "[[reddit.com-1]](https://www.reddit.com/r/LocalLLaMA/comments/1ilkosp/)"
            in result
        )
        assert (
            "[[reddit.com-2]](https://www.reddit.com/r/LocalLLaMA/comments/1j79obx/)"
            in result
        )
        assert (
            "[[youtube.com-1]](https://m.youtube.com/watch?v=4tDqV__jzKY)"
            in result
        )
        assert (
            "[[youtube.com-2]](https://www.youtube.com/watch?v=sGUjmyfof4Q)"
            in result
        )

        # Check single domain citations don't get IDs
        assert (
            "[[zilliz.com]](https://zilliz.com/blog/introduce-deepsearcher)"
            in result
        )
        assert (
            "[[milvus.io]](https://milvus.io/blog/introduce-deepsearcher.md)"
            in result
        )
        assert (
            "[[apidog.com]](https://apidog.com/blog/ollama-deep-research/)"
            in result
        )
        assert (
            "[[medium.com]](https://milvusio.medium.com/deepsearcher)" in result
        )

    def test_domain_id_always_hyperlinks_mode(self):
        """Test citation formatting with domain IDs always added for consistency."""
        # Use the same content as test_domain_id_hyperlinks_mode
        content_with_repeated_domains = """# Research Report

Citations from GitHub [1] and more from GitHub [2].
ArXiv papers [3] and another ArXiv paper [4].
A Nature article [5] and OpenAI research [6].

## Sources

[1] First GitHub Project
    URL: https://github.com/user/project1

[2] Second GitHub Project
    URL: https://github.com/user/project2

[3] First ArXiv Paper
    URL: https://arxiv.org/abs/2024.1111

[4] Second ArXiv Paper
    URL: https://arxiv.org/abs/2024.2222

[5] Nature Article
    URL: https://www.nature.com/articles/s41586-024-1234

[6] OpenAI Research
    URL: https://openai.com/research/gpt4
"""
        formatter = CitationFormatter(CitationMode.DOMAIN_ID_ALWAYS_HYPERLINKS)
        result = formatter.format_document(content_with_repeated_domains)

        # Check that ALL domains get IDs, even single ones
        assert "[[github.com-1]](https://github.com/user/project1)" in result
        assert "[[github.com-2]](https://github.com/user/project2)" in result
        assert "[[arxiv.org-1]](https://arxiv.org/abs/2024.1111)" in result
        assert "[[arxiv.org-2]](https://arxiv.org/abs/2024.2222)" in result
        assert (
            "[[nature.com-1]](https://www.nature.com/articles/s41586-024-1234)"
            in result
        )
        assert "[[openai.com-1]](https://openai.com/research/gpt4)" in result


class TestLaTeXExporter:
    """Test cases for LaTeX export functionality."""

    @pytest.fixture
    def sample_markdown(self):
        """Sample markdown for LaTeX conversion."""
        return """# Main Title

## Introduction

This is a **bold** statement with *italic* text and `code`.

### Subsection

- First item
- Second item
- Third item

Text with citations [1] and [2].

## Sources

[1] First Source
    URL: https://example.com
[2] Second Source
    URL: https://example2.com
"""

    def test_latex_export_basic(self, sample_markdown):
        """Test basic LaTeX export functionality."""
        exporter = LaTeXExporter()
        result = exporter.export_to_latex(sample_markdown)

        # Check header
        assert r"\documentclass[12pt]{article}" in result
        assert r"\begin{document}" in result
        assert r"\end{document}" in result

        # Check heading conversion
        assert r"\section{Main Title}" in result
        assert r"\subsection{Introduction}" in result
        assert r"\subsubsection{Subsection}" in result

        # Check emphasis conversion
        assert r"\textbf{bold}" in result
        assert r"\textit{italic}" in result
        assert r"\texttt{code}" in result

        # Check citation conversion
        assert r"\cite{1}" in result
        assert r"\cite{2}" in result

        # Check bibliography
        assert r"\begin{thebibliography}" in result
        assert r"\bibitem{1}" in result

    def test_latex_list_conversion(self, sample_markdown):
        """Test list conversion to LaTeX."""
        exporter = LaTeXExporter()
        result = exporter.export_to_latex(sample_markdown)

        # Check list conversion
        assert r"\begin{itemize}" in result
        assert r"\item First item" in result
        assert r"\item Second item" in result
        assert r"\item Third item" in result
        assert r"\end{itemize}" in result

    def test_latex_empty_content(self):
        """Test LaTeX export with empty content."""
        exporter = LaTeXExporter()
        result = exporter.export_to_latex("")

        assert r"\documentclass[12pt]{article}" in result
        assert r"\begin{document}" in result
        assert r"\end{document}" in result

    def test_latex_no_sources(self):
        """Test LaTeX export without sources section."""
        content = "# Title\n\nSome text without citations."
        exporter = LaTeXExporter()
        result = exporter.export_to_latex(content)

        # Should not include bibliography
        assert r"\begin{thebibliography}" not in result


class TestRISExporter:
    """Test cases for RIS export functionality."""

    @pytest.fixture
    def sample_markdown(self):
        """Sample markdown for RIS conversion."""
        return """# Research Report

This research explores various topics [1][2][3].

## Sources

[1] Understanding Deep Learning
    URL: https://arxiv.org/abs/2104.12345

[2] GitHub Project Documentation
    URL: https://github.com/example/project

[3] Reddit Discussion Thread
    URL: https://www.reddit.com/r/example/comments/abc123

[4] Local Resource Without URL
"""

    def test_ris_export_basic(self, sample_markdown):
        """Test basic RIS export functionality."""
        exporter = RISExporter()
        result = exporter.export_to_ris(sample_markdown)

        # Check that we have RIS entries
        assert "TY  - ELEC" in result
        assert "ER  - " in result

        # Check specific entries
        assert "ID  - ref1" in result
        assert "TI  - Understanding Deep Learning" in result
        assert "UR  - https://arxiv.org/abs/2104.12345" in result
        assert "PB  - arXiv" in result

        assert "ID  - ref2" in result
        assert "TI  - GitHub Project Documentation" in result
        assert "PB  - GitHub" in result

        assert "ID  - ref3" in result
        assert "TI  - Reddit Discussion Thread" in result
        assert "PB  - Reddit" in result

        # Check entry without URL
        assert "ID  - ref4" in result
        assert "TI  - Local Resource Without URL" in result

        # Check date fields
        assert "Y1  - " in result
        assert "DA  - " in result
        assert "LA  - en" in result

    def test_ris_export_empty_content(self):
        """Test RIS export with empty content."""
        exporter = RISExporter()
        result = exporter.export_to_ris("")
        assert result == ""

    def test_ris_export_no_sources(self):
        """Test RIS export without sources section."""
        content = "# Title\n\nSome text without sources."
        exporter = RISExporter()
        result = exporter.export_to_ris(content)
        assert result == ""

    def test_ris_publisher_extraction(self, sample_markdown):
        """Test that publishers are correctly extracted from domains."""
        exporter = RISExporter()
        result = exporter.export_to_ris(sample_markdown)

        # Check known publishers
        assert "PB  - arXiv" in result
        assert "PB  - GitHub" in result
        assert "PB  - Reddit" in result


class TestQuartoExporter:
    """Test cases for Quarto export functionality."""

    @pytest.fixture
    def sample_markdown(self):
        """Sample markdown for Quarto conversion."""
        return """# Research on AI Safety

## Introduction

This research explores **important** aspects of *AI safety* and `alignment`.

### Current Approaches

- Reinforcement learning from human feedback [1]
- Constitutional AI [2]
- Interpretability research [3]

Multiple studies [1, 2, 3] have shown promising results.

## Sources

[1] RLHF: A Survey
    URL: https://arxiv.org/abs/2023.1234
[2] Constitutional AI
    URL: https://anthropic.com/constitutional-ai
[3] Mechanistic Interpretability
    URL: https://distill.pub/2020/circuits
"""

    def test_quarto_export_basic(self, sample_markdown):
        """Test basic Quarto export functionality."""
        exporter = QuartoExporter()
        result = exporter.export_to_quarto(
            sample_markdown, "AI Safety Research"
        )

        # Check YAML header
        assert "---" in result
        assert 'title: "AI Safety Research"' in result
        assert 'author: "Local Deep Research"' in result
        # Check that date is in YYYY-MM-DD format
        import re

        assert re.search(r'date: "\d{4}-\d{2}-\d{2}"', result)
        assert "format:" in result
        assert "bibliography: references.bib" in result

        # Check citation conversion
        assert "[@ref1]" in result
        assert "[@ref2]" in result
        assert "[@ref3]" in result
        assert "[@ref1, @ref2, @ref3]" in result

        # Check bibliography note
        assert "Bibliography File Required" in result
        assert "references.bib" in result
        assert "@misc{ref1," in result

    def test_quarto_title_extraction(self, sample_markdown):
        """Test automatic title extraction from markdown."""
        exporter = QuartoExporter()
        # Don't provide title, should extract from markdown
        result = exporter.export_to_quarto(sample_markdown)

        assert 'title: "Research on AI Safety"' in result

    def test_quarto_bibliography_generation(self, sample_markdown):
        """Test BibTeX bibliography generation."""
        exporter = QuartoExporter()
        bib_content = exporter._generate_bibliography(sample_markdown)

        # Check BibTeX entries
        assert "@misc{ref1," in bib_content
        assert 'title = "{RLHF: A Survey}"' in bib_content
        assert "url = {https://arxiv.org/abs/2023.1234}" in bib_content
        assert (
            'howpublished = "\\url{https://arxiv.org/abs/2023.1234}"'
            in bib_content
        )
        assert "year = {2024}" in bib_content
        assert 'note = "Accessed: \\today"' in bib_content

        # Check all three entries
        assert "@misc{ref2," in bib_content
        assert "@misc{ref3," in bib_content

    def test_quarto_empty_content(self):
        """Test Quarto export with empty content."""
        exporter = QuartoExporter()
        result = exporter.export_to_quarto("", "Empty Report")

        # Should still have YAML header
        assert "---" in result
        assert 'title: "Empty Report"' in result

    def test_quarto_no_sources(self):
        """Test Quarto export without sources section."""
        content = "# Title\n\nSome text with citations [1] and [2]."
        exporter = QuartoExporter()
        result = exporter.export_to_quarto(content)

        # Citations should still be converted
        assert "[@ref1]" in result
        assert "[@ref2]" in result

        # Bibliography note should still appear but be empty
        assert "Bibliography File Required" in result
