import pytest
from docling_core.types.doc.document import DoclingDocument, TableData
from docling_core.types.doc.labels import DocItemLabel

from haiku.rag.client import HaikuRAG
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.client.processing import ensure_chunks_embedded
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult


async def create_document_with_docling(
    client: HaikuRAG, docling_doc: DoclingDocument, title: str
):
    """Helper to create a document from a DoclingDocument using import_document."""
    chunks = await client.chunk(docling_doc)
    embedded_chunks = await ensure_chunks_embedded(client._config, chunks)
    return await client.import_document(
        docling_document=docling_doc,
        chunks=embedded_chunks,
        title=title,
    )


def create_table_document() -> DoclingDocument:
    """Create a document with a table that will be split across chunks."""
    doc = DoclingDocument(name="table_test")

    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Introduction paragraph.")
    doc.add_heading(text="Employee Data", level=1)

    # Create a table with enough content to span multiple chunks
    table_data = TableData(num_cols=3, num_rows=0)
    table_data.add_row(["Name", "Age", "City"])
    table_data.add_row(["Alice Smith", "30", "New York"])
    table_data.add_row(["Bob Johnson", "25", "Los Angeles"])
    table_data.add_row(["Charlie Brown", "35", "Chicago"])
    table_data.add_row(["Diana Ross", "28", "Miami"])
    doc.add_table(data=table_data)

    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Conclusion paragraph.")

    return doc


def create_list_document() -> DoclingDocument:
    """Create a document with list items that will be split across chunks."""
    doc = DoclingDocument(name="list_test")

    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Shopping list for the week:")
    doc.add_list_item(
        text="Fresh organic apples from the farmers market", enumerated=False
    )
    doc.add_list_item(text="Ripe yellow bananas for smoothies", enumerated=False)
    doc.add_list_item(text="Valencia oranges for fresh juice", enumerated=False)
    doc.add_list_item(text="Seedless red grapes as healthy snack", enumerated=False)
    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Remember to bring reusable bags.")

    return doc


def create_code_document() -> DoclingDocument:
    """Create a document with adjacent code blocks that will be split."""
    doc = DoclingDocument(name="code_test")

    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Here are several code snippets:")
    # Multiple adjacent code blocks - type-aware expansion should group them
    doc.add_text(label=DocItemLabel.CODE, text="# Part 1: Setup\nimport os\nimport sys")
    doc.add_text(
        label=DocItemLabel.CODE, text='# Part 2: Config\nCONFIG = {"debug": True}'
    )
    doc.add_text(
        label=DocItemLabel.CODE, text="# Part 3: Main\ndef main():\n    print(CONFIG)"
    )
    doc.add_text(label=DocItemLabel.PARAGRAPH, text="End of code examples.")

    return doc


@pytest.fixture
def small_chunk_config() -> AppConfig:
    """Config with small chunk size to force splitting."""
    config = AppConfig()
    config.processing.chunk_size = 32
    config.search.max_context_chars = 10000
    return config


@pytest.mark.vcr()
async def test_table_expansion_includes_split_rows(temp_db_path, small_chunk_config):
    """Verify that table expansion retrieves rows that were split into different chunks."""
    docling_doc = create_table_document()

    async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
        doc = await create_document_with_docling(client, docling_doc, "Table Test")
        assert doc.id is not None

        # Search for table content
        results = await client.search("Alice Smith New York employee", limit=5)
        table_results = [r for r in results if "table" in r.labels]
        assert len(table_results) > 0, (
            f"No table results. Labels: {[r.labels for r in results]}"
        )

        original = table_results[0]

        # Verify the original chunk does NOT contain all table data
        # (proving we need expansion)
        original_has_all = all(
            name in original.content for name in ["Alice", "Bob", "Charlie", "Diana"]
        )

        # Expand context
        expanded = await client.expand_context(table_results[:1])
        assert len(expanded) == 1

        expanded_content = expanded[0].content

        # After expansion, we should have the complete table
        assert "Alice" in expanded_content
        assert "Bob" in expanded_content
        assert "Charlie" in expanded_content
        assert "Diana" in expanded_content

        # Verify expansion actually added content (unless chunk already had everything)
        if not original_has_all:
            assert len(expanded_content) > len(original.content), (
                "Expansion should have added content"
            )


@pytest.mark.vcr()
async def test_list_expansion_includes_split_items(temp_db_path, small_chunk_config):
    """Verify that list expansion retrieves items that were split into different chunks."""
    docling_doc = create_list_document()

    async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
        doc = await create_document_with_docling(client, docling_doc, "List Test")
        assert doc.id is not None

        # Search for a list item
        results = await client.search("grapes healthy snack", limit=5)
        list_results = [r for r in results if "list_item" in r.labels]
        assert len(list_results) > 0, (
            f"No list results. Labels: {[r.labels for r in results]}"
        )

        original = list_results[0]

        # Check what the original chunk contains
        original_items = sum(
            1
            for item in ["apples", "bananas", "oranges", "grapes"]
            if item in original.content.lower()
        )

        # Expand context
        expanded = await client.expand_context(list_results[:1])
        assert len(expanded) == 1

        expanded_content = expanded[0].content.lower()

        # Count items after expansion
        expanded_items = sum(
            1
            for item in ["apples", "bananas", "oranges", "grapes"]
            if item in expanded_content
        )

        # Expansion should include at least as many items (more if split)
        assert expanded_items >= original_items

        # If original didn't have all items, expansion should have added some
        if original_items < 4:
            assert expanded_items > original_items, (
                f"Expansion should have added items. Original: {original_items}, Expanded: {expanded_items}"
            )


@pytest.mark.vcr()
async def test_code_expansion_includes_adjacent_blocks(
    temp_db_path, small_chunk_config
):
    """Verify that code expansion retrieves adjacent code blocks split across chunks."""
    docling_doc = create_code_document()

    async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
        doc = await create_document_with_docling(client, docling_doc, "Code Test")
        assert doc.id is not None

        # Search for middle code block (Part 2)
        results = await client.search("CONFIG debug True", limit=5)
        code_results = [r for r in results if "code" in r.labels]
        assert len(code_results) > 0, (
            f"No code results. Labels: {[r.labels for r in results]}"
        )

        original = code_results[0]

        # Check what parts the original chunk has
        original_parts = sum(
            1 for part in ["Part 1", "Part 2", "Part 3"] if part in original.content
        )

        # Expand context
        expanded = await client.expand_context(code_results[:1])
        assert len(expanded) == 1

        expanded_content = expanded[0].content

        # Count parts after expansion
        expanded_parts = sum(
            1 for part in ["Part 1", "Part 2", "Part 3"] if part in expanded_content
        )

        # Expansion should include at least as many parts
        assert expanded_parts >= original_parts

        # If original didn't have all parts, expansion should have added some
        if original_parts < 3:
            assert expanded_parts > original_parts, (
                f"Expansion should have added code blocks. Original: {original_parts}, Expanded: {expanded_parts}"
            )


@pytest.mark.vcr()
async def test_text_expansion_includes_surrounding(temp_db_path):
    """Text content expansion should include surrounding paragraphs."""
    config = AppConfig()
    config.processing.chunk_size = 32

    # Create a document with longer paragraphs that will split
    doc = DoclingDocument(name="text_test")
    doc.add_text(
        label=DocItemLabel.PARAGRAPH,
        text="First paragraph with enough content to be its own chunk in the document.",
    )
    doc.add_text(
        label=DocItemLabel.PARAGRAPH,
        text="Second paragraph contains different information about software testing.",
    )
    doc.add_text(
        label=DocItemLabel.PARAGRAPH,
        text="Third paragraph discusses various topics and provides more details.",
    )
    doc.add_text(
        label=DocItemLabel.PARAGRAPH,
        text="Fourth paragraph concludes the document with final thoughts.",
    )

    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
        document = await create_document_with_docling(client, doc, "Text Only")
        assert document.id is not None

        # Search for second paragraph
        results = await client.search("software testing", limit=1)
        assert len(results) > 0

        original = results[0]
        expanded = await client.expand_context(results)

        # Expansion should include adjacent paragraphs
        assert len(expanded[0].content) >= len(original.content)


@pytest.mark.vcr()
async def test_expansion_preserves_metadata(temp_db_path, small_chunk_config):
    """Expansion should preserve document metadata."""
    docling_doc = create_table_document()

    async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
        doc = await create_document_with_docling(client, docling_doc, "Metadata Test")
        assert doc.id is not None

        results = await client.search("Introduction paragraph", limit=1)
        assert len(results) > 0

        expanded = await client.expand_context(results)

        assert expanded[0].document_title == "Metadata Test"
        assert expanded[0].chunk_id == results[0].chunk_id
        assert expanded[0].document_id == results[0].document_id


@pytest.mark.vcr()
async def test_format_for_agent_output(temp_db_path, small_chunk_config):
    """format_for_agent should include source, type, and content sections."""
    docling_doc = create_table_document()

    async with HaikuRAG(temp_db_path, config=small_chunk_config, create=True) as client:
        doc = await create_document_with_docling(client, docling_doc, "Format Test")
        assert doc.id is not None

        results = await client.search("Alice Smith employee data", limit=5)
        table_results = [r for r in results if "table" in r.labels]
        assert len(table_results) > 0

        expanded = await client.expand_context(table_results[:1])
        # Format with rank (the way agents use it)
        formatted = expanded[0].format_for_agent(rank=1, total=1)

        # Check format structure
        assert "[rank 1 of 1]" in formatted
        assert 'Source: "Format Test"' in formatted
        assert "Type: table" in formatted
        assert "Content:" in formatted


async def test_expand_context_single_item_document(temp_db_path):
    """Test expand_context with a single-item document."""
    from haiku.rag.store.models.document import Document

    docling_doc = DoclingDocument(name="simple")
    docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Simple test content")

    async with HaikuRAG(temp_db_path, create=True) as client:
        document = Document(content="Simple test content")
        document.set_docling(docling_doc)
        doc = await _store_document_with_chunks(client, document, [], docling_doc)
        assert doc.id is not None

        # Create a search result with a doc_item_ref pointing to the item
        items = await client.document_item_repository.get_items_in_range(doc.id, 0, 10)
        assert len(items) > 0

        search_results = [
            SearchResult(
                content="Simple test content",
                score=0.9,
                document_id=doc.id,
                doc_item_refs=[items[0].self_ref],
            )
        ]
        expanded_results = await client.expand_context(search_results)

        assert len(expanded_results) == 1
        assert expanded_results[0].score == 0.9
        assert "Simple test content" in expanded_results[0].content


async def test_expand_context_no_refs_passes_through(temp_db_path):
    """Results without doc_item_refs pass through unexpanded."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # A search result with no doc_item_refs should pass through as-is
        search_results = [
            SearchResult(
                content="Some chunk content",
                score=0.8,
                document_id="some-doc",
                doc_item_refs=[],
            )
        ]
        expanded = await client.expand_context(search_results)

        assert len(expanded) == 1
        assert expanded[0].content == "Some chunk content"
        assert expanded[0].score == 0.8


@pytest.mark.vcr()
async def test_expand_context_with_docling_merges_overlapping(temp_db_path):
    """Test that expand_context with DoclingDocument merges overlapping results."""
    config = AppConfig()

    markdown_content = """# Chapter 1

This is paragraph one about topic A.

This is paragraph two about topic A continued.

This is paragraph three about topic B.

# Chapter 2

This is paragraph four about topic C.
"""

    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
        doc = await client.create_document(
            content=markdown_content,
            uri="test://structured",
        )

        assert doc.id is not None
        assert doc.docling_document is not None

        # Get chunks which should have doc_item_refs
        chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(chunks) >= 1

        # Find chunks that have doc_item_refs (from docling chunking)
        chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]

        if len(chunks_with_refs) >= 2:
            # Create search results from adjacent chunks
            search_results = [
                SearchResult.from_chunk(chunks_with_refs[0], 0.9),
                SearchResult.from_chunk(chunks_with_refs[1], 0.8),
            ]

            # Expand with configured radius that should cause overlap
            expanded = await client.expand_context(search_results)

            # If chunks were adjacent, they should be merged
            # The expanded results should have merged metadata
            assert len(expanded) >= 1

            # Check that expanded result has page_numbers populated
            for r in expanded:
                # Should have doc_item_refs from expansion
                assert r.doc_item_refs is not None


@pytest.mark.vcr()
async def test_expand_context_docling_merges_metadata(temp_db_path):
    """Test that expand_context properly merges metadata from multiple results."""
    config = AppConfig()

    markdown_content = """# Introduction

First paragraph of introduction.

Second paragraph of introduction.

# Methods

First paragraph of methods section.

Second paragraph of methods section.

# Results

First paragraph of results.
"""

    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
        doc = await client.create_document(
            content=markdown_content,
            uri="test://metadata-merge",
        )

        assert doc.id is not None

        chunks = await client.chunk_repository.get_by_document_id(doc.id)
        chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]

        if len(chunks_with_refs) >= 2:
            # Get chunks with different headings if possible
            chunk1 = chunks_with_refs[0]
            chunk2 = chunks_with_refs[-1]  # Last chunk likely has different heading

            search_results = [
                SearchResult.from_chunk(chunk1, 0.9),
                SearchResult.from_chunk(chunk2, 0.8),
            ]

            # Expand with large radius to potentially merge
            expanded = await client.expand_context(search_results)

            # Check that results have proper structure
            for r in expanded:
                # Content should be non-empty
                assert len(r.content) > 0
                # Score should be preserved (best score)
                assert r.score in [0.9, 0.8]
                # Expanded content should have docling refs
                assert r.doc_item_refs is not None and len(r.doc_item_refs) > 0
                # Document has headings, expanded result should too
                assert r.headings is not None and len(r.headings) > 0


def create_picture_document() -> DoclingDocument:
    """Create a document with picture items for testing image handling."""
    from docling_core.types.doc.labels import DocItemLabel

    doc = DoclingDocument(name="picture_test")

    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Introduction with context.")
    doc.add_heading(text="Figures Section", level=1)

    # Add a picture item - this will have export_to_markdown method
    # that defaults to EMBEDDED mode if called without image_mode parameter
    doc.add_picture()

    doc.add_text(label=DocItemLabel.CAPTION, text="Figure 1: Sample diagram")
    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Conclusion paragraph.")

    return doc


@pytest.mark.vcr()
async def test_expand_context_no_base64_images(temp_db_path):
    """Ensure expanded context does not contain base64 image data.

    This test verifies that when expand_context includes PictureItem objects,
    they are serialized with PLACEHOLDER mode (not EMBEDDED), preventing
    base64 image data from leaking into the expanded content.
    """
    config = AppConfig()

    docling_doc = create_picture_document()

    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
        doc = await create_document_with_docling(client, docling_doc, "Picture Test")
        assert doc.id is not None

        # Search for content near the picture
        results = await client.search("Figure diagram", limit=5)
        assert len(results) > 0

        # Expand context - this should include the picture area
        expanded = await client.expand_context(results)

        for result in expanded:
            # Base64 image data should never appear in expanded content
            assert "base64" not in result.content.lower(), (
                f"Found 'base64' in expanded content: {result.content[:500]}"
            )
            assert "data:image" not in result.content.lower(), (
                f"Found 'data:image' in expanded content: {result.content[:500]}"
            )


@pytest.mark.vcr()
async def test_expand_context_no_base64_images_docling_local(
    temp_db_path, doclaynet_first_page_pdf
):
    """Ensure expanded context from real PDF does not contain base64 image data.

    Tests end-to-end with a single page of doclaynet.pdf using docling-local converter.
    """
    config = AppConfig()
    config.processing.converter = "docling-local"
    config.processing.chunker = "docling-local"
    config.processing.conversion_options.do_ocr = False

    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
        pdf_path = doclaynet_first_page_pdf
        result = await client.create_document_from_source(pdf_path)
        doc = result if not isinstance(result, list) else result[0]
        assert doc.id is not None

        # Search for content that might include pictures
        results = await client.search("figure table", limit=10)

        if len(results) > 0:
            expanded = await client.expand_context(results)

            for result in expanded:
                assert "base64" not in result.content.lower(), (
                    f"Found 'base64' in expanded content: {result.content[:500]}"
                )
                assert "data:image" not in result.content.lower(), (
                    f"Found 'data:image' in expanded content: {result.content[:500]}"
                )


@pytest.mark.vcr()
async def test_expand_context_no_base64_images_docling_serve(temp_db_path):
    """Ensure expanded context from real PDF does not contain base64 image data.

    Tests end-to-end with doclaynet.pdf using docling-serve converter.
    """
    from pathlib import Path

    config = AppConfig()
    config.processing.converter = "docling-serve"
    config.processing.chunker = "docling-serve"

    async with HaikuRAG(temp_db_path, config=config, create=True) as client:
        pdf_path = Path(__file__).parent / "data" / "doclaynet.pdf"
        result = await client.create_document_from_source(pdf_path)
        doc = result if not isinstance(result, list) else result[0]
        assert doc.id is not None

        # Search for content that might include pictures
        results = await client.search("figure table", limit=10)

        if len(results) > 0:
            expanded = await client.expand_context(results)

            for result in expanded:
                assert "base64" not in result.content.lower(), (
                    f"Found 'base64' in expanded content: {result.content[:500]}"
                )
                assert "data:image" not in result.content.lower(), (
                    f"Found 'data:image' in expanded content: {result.content[:500]}"
                )