haiku.rag/tests/test_client.py

import tempfile
from pathlib import Path
from unittest.mock import AsyncMock, patch

import httpx
import pytest
from datasets import Dataset

from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document


@pytest.mark.asyncio
async def test_client_document_crud(qa_corpus: Dataset, temp_db_path):
    """Test HaikuRAG CRUD operations for documents."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Get test data
        first_doc = qa_corpus[0]
        document_text = first_doc["document_extracted"]
        test_uri = "file:///path/to/test.txt"
        test_metadata = {"source": "test", "topic": "testing"}

        # Test create_document
        created_doc = await client.create_document(
            content=document_text, uri=test_uri, metadata=test_metadata
        )

        assert created_doc.id is not None
        assert created_doc.content == document_text
        assert created_doc.uri == test_uri
        assert created_doc.metadata == test_metadata

        # Test get_document_by_id
        retrieved_doc = await client.get_document_by_id(created_doc.id)
        assert retrieved_doc is not None
        assert retrieved_doc.id == created_doc.id
        assert retrieved_doc.content == document_text
        assert retrieved_doc.uri == test_uri

        # Test get_document_by_uri
        retrieved_by_uri = await client.get_document_by_uri(test_uri)
        assert retrieved_by_uri is not None
        assert retrieved_by_uri.id == created_doc.id
        assert retrieved_by_uri.content == document_text

        # Test get_document_by_uri with non-existent URI
        non_existent = await client.get_document_by_uri("file:///non/existent.txt")
        assert non_existent is None

        # Test update_document
        retrieved_doc.content = "Updated content"
        retrieved_doc.uri = "file:///updated/path.txt"
        updated_doc = await client.update_document(retrieved_doc)
        assert updated_doc.content == "Updated content"
        assert updated_doc.uri == "file:///updated/path.txt"

        # Test list_documents
        all_docs = await client.list_documents()
        assert len(all_docs) == 1
        assert all_docs[0].id == created_doc.id

        # Test list_documents with pagination
        limited_docs = await client.list_documents(limit=10, offset=0)
        assert len(limited_docs) == 1

        # Test delete_document
        deleted = await client.delete_document(created_doc.id)
        assert deleted is True

        # Verify document is gone
        retrieved_doc = await client.get_document_by_id(created_doc.id)
        assert retrieved_doc is None

        # Test delete non-existent document
        deleted_again = await client.delete_document(created_doc.id)
        assert deleted_again is False


@pytest.mark.asyncio
async def test_client_update_document_fields(qa_corpus: Dataset, temp_db_path):
    """Test updating document with individual parameters."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Get test data
        first_doc = qa_corpus[0]
        document_text = first_doc["document_extracted"]
        test_uri = "file:///path/to/test.txt"
        test_metadata = {"source": "test", "topic": "testing"}

        # Create a document
        created_doc = await client.create_document(
            content=document_text,
            uri=test_uri,
            title="Original Title",
            metadata=test_metadata,
        )
        assert created_doc.id is not None
        original_id = created_doc.id

        # Test updating only content
        updated_doc = await client.update_document_fields(
            document_id=original_id, content="Updated content only"
        )
        assert updated_doc.id == original_id
        assert updated_doc.content == "Updated content only"
        assert updated_doc.title == "Original Title"
        assert updated_doc.uri == test_uri

        # Test updating only metadata
        new_metadata = {"source": "updated", "version": "2.0"}
        updated_doc = await client.update_document_fields(
            document_id=original_id, metadata=new_metadata
        )
        assert updated_doc.metadata == new_metadata
        assert (
            updated_doc.content == "Updated content only"
        )  # Should keep previous update

        # Test updating only title
        updated_doc = await client.update_document_fields(
            document_id=original_id, title="New Title"
        )
        assert updated_doc.title == "New Title"
        assert updated_doc.content == "Updated content only"
        assert updated_doc.metadata == new_metadata

        # Test updating multiple fields at once
        custom_chunks = [
            Chunk(content="Custom chunk 1", order=0),
            Chunk(content="Custom chunk 2", order=1),
        ]
        updated_doc = await client.update_document_fields(
            document_id=original_id,
            content="Content with custom chunks",
            title="Final Title",
            metadata={"final": "true"},
            chunks=custom_chunks,
        )
        assert updated_doc.id == original_id
        assert updated_doc.content == "Content with custom chunks"
        assert updated_doc.title == "Final Title"
        assert updated_doc.metadata == {"final": "true"}

        # Verify the custom chunks were created
        doc_chunks = await client.chunk_repository.get_by_document_id(original_id)
        assert len(doc_chunks) == 2
        assert doc_chunks[0].content == "Custom chunk 1"
        assert doc_chunks[1].content == "Custom chunk 2"


@pytest.mark.asyncio
async def test_client_create_document_from_source(temp_db_path):
    """Test creating a document from a file source."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            test_content = "This is test content from a file."
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text(test_content)

            # Test create_document_from_source with Path
            doc = await client.create_document_from_source(source=temp_path)
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert doc.content == test_content
            assert doc.uri == temp_path.as_uri()
            assert "contentType" in doc.metadata
            assert "md5" in doc.metadata
            assert doc.metadata["contentType"] == "text/plain"

            # Test create_document_from_source with string path
            doc2 = await client.create_document_from_source(source=str(temp_path))
            assert isinstance(doc2, Document)

            assert doc2.id is not None
            assert doc2.content == test_content
            assert doc2.uri == temp_path.as_uri()
            assert "contentType" in doc2.metadata
            assert "md5" in doc2.metadata


@pytest.mark.asyncio
async def test_client_create_document_from_source_with_title(temp_db_path):
    """Test creating a document from a file source with a title."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            test_content = "This is test content from a file."
            temp_path = Path(temp_dir) / "test_title.txt"
            temp_path.write_text(test_content)

            doc = await client.create_document_from_source(
                source=temp_path, title="My Doc"
            )
            assert isinstance(doc, Document)
            assert doc.id is not None
            assert doc.title == "My Doc"


@pytest.mark.asyncio
async def test_client_update_title_noop_behavior(temp_db_path):
    """When content is unchanged, updating title should update document without re-chunking."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test_update_title.txt"
            temp_path.write_text("Original content")

            doc1 = await client.create_document_from_source(temp_path, title="Title A")
            assert isinstance(doc1, Document)
            assert doc1.id is not None

            # Re-add with same content but new title
            doc2 = await client.create_document_from_source(temp_path, title="Title B")
            assert isinstance(doc2, Document)
            assert doc2.id == doc1.id
            # Fetch and verify title updated
            got = await client.get_document_by_id(doc1.id)
            assert got is not None
            assert got.title == "Title B"


@pytest.mark.asyncio
async def test_client_create_document_from_source_unsupported(temp_db_path):
    """Test creating a document from an unsupported file type."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a temporary file with unsupported extension
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".unsupported", delete=False
        ) as f:
            f.write("content")
            temp_path = Path(f.name)

            # Should raise ValueError for unsupported extension
            with pytest.raises(ValueError, match="Unsupported file extension"):
                await client.create_document_from_source(temp_path)


@pytest.mark.asyncio
async def test_client_create_document_from_source_nonexistent(temp_db_path):
    """Test creating a document from a non-existent file."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        non_existent_path = Path("/non/existent/file.txt")

        # Should raise ValueError when file doesn't exist
        with pytest.raises(ValueError, match="File does not exist"):
            await client.create_document_from_source(non_existent_path)


@pytest.mark.asyncio
async def test_client_create_document_from_directory(temp_db_path):
    """Test creating documents from a directory recursively."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            test_dir = Path(temp_dir) / "test_docs"
            test_dir.mkdir()

            (test_dir / "doc1.txt").write_text("Content of doc1")
            (test_dir / "doc2.md").write_text("# Content of doc2")

            subdir = test_dir / "subdir"
            subdir.mkdir()
            (subdir / "doc3.py").write_text("print('hello')")

            (test_dir / "unsupported.xyz").write_text("unsupported file")

            result = await client.create_document_from_source(test_dir)

            assert isinstance(result, list)
            assert len(result) == 3

            for doc in result:
                assert doc.id is not None
                assert doc.uri is not None
                assert "md5" in doc.metadata
                assert "contentType" in doc.metadata

            uris = [doc.uri for doc in result if doc.uri]
            assert any("doc1.txt" in uri for uri in uris)
            assert any("doc2.md" in uri for uri in uris)
            assert any("doc3.py" in uri for uri in uris)
            assert not any("unsupported.xyz" in uri for uri in uris)


@pytest.mark.asyncio
async def test_client_create_document_from_directory_with_filters(
    monkeypatch, temp_db_path
):
    """Test creating documents from a directory with ignore and include patterns."""
    # Mock config to have ignore and include patterns
    monkeypatch.setattr(
        "haiku.rag.client.Config.monitor.ignore_patterns", ["**/ignore_me/**", "*.log"]
    )
    monkeypatch.setattr(
        "haiku.rag.client.Config.monitor.include_patterns", ["**/include/**/*.txt"]
    )

    async with HaikuRAG(temp_db_path, create=True) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            test_dir = Path(temp_dir) / "test_docs"
            test_dir.mkdir()

            # Create files in include directory - should be included
            include_dir = test_dir / "include"
            include_dir.mkdir()
            (include_dir / "doc1.txt").write_text("Content of doc1")
            (include_dir / "doc2.txt").write_text("Content of doc2")

            # Create files outside include directory - should be excluded by include pattern
            (test_dir / "doc3.txt").write_text("Content of doc3")

            # Create files in ignore directory - should be excluded by ignore pattern
            ignore_dir = test_dir / "ignore_me"
            ignore_dir.mkdir()
            (ignore_dir / "doc4.txt").write_text("Content of doc4")

            # Create log file - should be excluded by ignore pattern
            (test_dir / "debug.log").write_text("log content")

            result = await client.create_document_from_source(test_dir)

            assert isinstance(result, list)
            # Should only include doc1.txt and doc2.txt from include directory
            assert len(result) == 2

            uris = [doc.uri for doc in result if doc.uri]
            assert any("doc1.txt" in uri for uri in uris)
            assert any("doc2.txt" in uri for uri in uris)
            assert not any("doc3.txt" in uri for uri in uris)
            assert not any("doc4.txt" in uri for uri in uris)
            assert not any("debug.log" in uri for uri in uris)


@pytest.mark.asyncio
async def test_client_create_document_from_url(temp_db_path):
    """Test creating a document from a URL."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Mock the HTTP response
        mock_response = AsyncMock()
        mock_response.content = b"<html><body><h1>Test Page</h1><p>This is test content from a webpage.</p></body></html>"
        mock_response.headers = {"content-type": "text/html"}
        mock_response.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_response):
            doc = await client.create_document_from_source(
                source="https://example.com/test.html", metadata={"source_type": "web"}
            )
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert "Test Page" in doc.content
            assert "test content" in doc.content
            assert doc.uri == "https://example.com/test.html"
            assert doc.metadata["source_type"] == "web"
            assert "contentType" in doc.metadata
            assert "md5" in doc.metadata
            assert doc.metadata["contentType"] == "text/html"


@pytest.mark.asyncio
async def test_client_create_document_from_url_with_different_content_types(
    temp_db_path,
):
    """Test creating documents from URLs with different content types."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Test JSON content
        mock_json_response = AsyncMock()
        mock_json_response.content = (
            b'{"title": "Test JSON", "content": "This is JSON content"}'
        )
        mock_json_response.headers = {"content-type": "application/json"}
        mock_json_response.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_json_response):
            doc = await client.create_document_from_source(
                "https://api.example.com/data.json"
            )
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert "Test JSON" in doc.content
            assert doc.uri == "https://api.example.com/data.json"
            assert "contentType" in doc.metadata
            assert "md5" in doc.metadata
            assert doc.metadata["contentType"] == "application/json"

        # Test plain text content
        mock_text_response = AsyncMock()
        mock_text_response.content = b"This is plain text content from a URL."
        mock_text_response.headers = {"content-type": "text/plain"}
        mock_text_response.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_text_response):
            doc = await client.create_document_from_source(
                "https://example.com/readme.txt"
            )
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert doc.content == "This is plain text content from a URL."
            assert doc.uri == "https://example.com/readme.txt"
            assert "contentType" in doc.metadata
            assert "md5" in doc.metadata
            assert doc.metadata["contentType"] == "text/plain"


@pytest.mark.asyncio
async def test_client_create_document_from_url_unsupported_content(temp_db_path):
    """Test creating a document from URL with unsupported content type."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Mock response with unsupported content type
        mock_response = AsyncMock()
        mock_response.content = b"binary content"
        mock_response.headers = {"content-type": "application/octet-stream"}
        mock_response.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_response):
            with pytest.raises(ValueError, match="Unsupported content type"):
                await client.create_document_from_source(
                    "https://example.com/binary.bin"
                )


@pytest.mark.asyncio
async def test_client_create_document_from_url_http_error(temp_db_path):
    """Test handling HTTP errors when creating document from URL."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with patch("httpx.AsyncClient.get") as mock_get:
            mock_get.side_effect = httpx.HTTPStatusError(
                "404 Not Found",
                request=httpx.Request("GET", "https://example.com/notfound.html"),
                response=httpx.Response(404),
            )

            with pytest.raises(httpx.HTTPStatusError):
                await client.create_document_from_source(
                    "https://example.com/notfound.html"
                )


@pytest.mark.asyncio
async def test_get_extension_from_content_type_or_url(temp_db_path):
    """Test the helper method for determining file extensions."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Test content type mappings
        assert (
            client._get_extension_from_content_type_or_url("", "text/html") == ".html"
        )
        assert (
            client._get_extension_from_content_type_or_url("", "application/pdf")
            == ".pdf"
        )
        assert (
            client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
        )

        # Test URL extension detection
        assert (
            client._get_extension_from_content_type_or_url(
                "https://example.com/doc.pdf", ""
            )
            == ".pdf"
        )
        assert (
            client._get_extension_from_content_type_or_url(
                "https://example.com/data.json", ""
            )
            == ".json"
        )

        # Test default fallback
        assert (
            client._get_extension_from_content_type_or_url("https://example.com/", "")
            == ".html"
        )

        # Test content type priority over URL extension
        assert (
            client._get_extension_from_content_type_or_url(
                "https://example.com/file.txt", "application/pdf"
            )
            == ".pdf"
        )


@pytest.mark.asyncio
async def test_client_metadata_content_type_and_md5(temp_db_path):
    """Test that contentType and md5 metadata are correctly set."""
    import hashlib

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a temporary file with known content
        test_content = "Test content for MD5 calculation."
        expected_md5 = hashlib.md5(test_content.encode()).hexdigest()

        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text(test_content)

            doc = await client.create_document_from_source(temp_path)
            assert isinstance(doc, Document)

            assert doc.metadata["contentType"] == "text/plain"
            assert doc.metadata["md5"] == expected_md5

            mock_response = AsyncMock()
            mock_response.content = test_content.encode()
            mock_response.headers = {"content-type": "text/plain"}
            mock_response.raise_for_status = AsyncMock()

            with patch("httpx.AsyncClient.get", return_value=mock_response):
                url_doc = await client.create_document_from_source(
                    "https://example.com/test.txt"
                )
                assert isinstance(url_doc, Document)

                assert url_doc.metadata["contentType"] == "text/plain"
                assert url_doc.metadata["md5"] == expected_md5


@pytest.mark.asyncio
async def test_client_create_update_no_op_behavior(temp_db_path):
    """Test create/update/no-op behavior based on MD5 changes."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a temporary file
        test_content = "Original content for testing."
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text(test_content)

            # First call - should create new document
            doc1 = await client.create_document_from_source(temp_path)
            assert isinstance(doc1, Document)
            assert doc1.id is not None
            assert doc1.content == test_content
            original_id = doc1.id

            # Second call with same content - should return existing document (no-op)
            doc2 = await client.create_document_from_source(temp_path)
            assert isinstance(doc2, Document)
            assert doc2.id == original_id  # Same document
            assert doc2.content == test_content

            # Modify file content
            updated_content = "Updated content for testing."
            temp_path.write_text(updated_content)

            # Third call with changed content - should update existing document
            doc3 = await client.create_document_from_source(temp_path)
            assert isinstance(doc3, Document)
            assert doc3.id == original_id  # Same document ID
            assert doc3.content == updated_content  # Updated content

            # Verify the document was actually updated in database
            retrieved_doc = await client.get_document_by_id(original_id)
            assert retrieved_doc is not None
            assert retrieved_doc.content == updated_content


@pytest.mark.asyncio
async def test_client_unchanged_file_keeps_timestamp(temp_db_path):
    """Test that unchanged files don't update the updated_at timestamp."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a temporary file
        test_content = "Test content for timestamp check."
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text(test_content)

            # First call - create document
            doc1 = await client.create_document_from_source(temp_path)
            assert isinstance(doc1, Document)
            original_updated_at = doc1.updated_at

            # Second call with same content - should not update timestamp
            doc2 = await client.create_document_from_source(temp_path)
            assert isinstance(doc2, Document)
            assert doc2.id == doc1.id
            assert doc2.updated_at == original_updated_at  # Timestamp should not change


@pytest.mark.asyncio
async def test_client_url_create_update_no_op_behavior(temp_db_path):
    """Test create/update/no-op behavior for URLs based on MD5 changes."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        url = "https://example.com/test.txt"
        original_content = b"Original URL content"
        updated_content = b"Updated URL content"

        # Mock first response
        mock_response1 = AsyncMock()
        mock_response1.content = original_content
        mock_response1.headers = {"content-type": "text/plain"}
        mock_response1.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_response1):
            # First call - should create new document
            doc1 = await client.create_document_from_source(url)
            assert isinstance(doc1, Document)
            assert doc1.id is not None
            original_id = doc1.id

            # Second call with same content - should return existing document (no-op)
            doc2 = await client.create_document_from_source(url)
            assert isinstance(doc2, Document)
            assert doc2.id == original_id  # Same document

        mock_response2 = AsyncMock()
        mock_response2.content = updated_content
        mock_response2.headers = {"content-type": "text/plain"}
        mock_response2.raise_for_status = AsyncMock()

        with patch("httpx.AsyncClient.get", return_value=mock_response2):
            # Third call with changed content - should update existing document
            doc3 = await client.create_document_from_source(url)
            assert isinstance(doc3, Document)
            assert doc3.id == original_id  # Same document ID
            assert doc3.content == updated_content.decode()  # Updated content


@pytest.mark.asyncio
async def test_client_search(temp_db_path):
    """Test HaikuRAG search functionality."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Add multiple documents to search from
        doc1_text = "Python is a high-level programming language known for its simplicity and readability."
        doc2_text = "Machine learning algorithms help computers learn patterns from data without explicit programming."
        doc3_text = "Data science combines statistics, programming, and domain expertise to extract insights."

        # Create documents
        doc1 = await client.create_document(
            content=doc1_text, uri="doc1.txt", metadata={"topic": "python"}
        )
        doc2 = await client.create_document(
            content=doc2_text, uri="doc2.txt", metadata={"topic": "ml"}
        )
        await client.create_document(
            content=doc3_text, uri="doc3.txt", metadata={"topic": "data_science"}
        )

        # Test search with keyword that should match doc1
        results = await client.search("Python programming", limit=3)

        assert len(results) > 0
        # Verify results are SearchResult objects with expected fields
        first_result = results[0]
        assert first_result.content
        assert first_result.score >= 0
        assert first_result.document_id == doc1.id

        # Test search with different query
        ml_results = await client.search("machine learning data", limit=2)
        assert len(ml_results) > 0

        # Verify first result is from the machine learning document (doc2)
        first_ml_result = ml_results[0]
        assert first_ml_result.document_id == doc2.id

        # Test search with limit parameter
        limited_results = await client.search("programming", limit=1)
        assert len(limited_results) <= 1


@pytest.mark.asyncio
async def test_client_async_context_manager(temp_db_path):
    """Test HaikuRAG as async context manager."""

    # Test that context manager works and auto-closes
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a document to ensure the client works
        doc = await client.create_document(
            content="Test content for context manager",
            uri="test://context",
            metadata={"test": "context_manager"},
        )

        assert doc.id is not None
        assert doc.content == "Test content for context manager"

        # Test search works within context
        results = await client.search("Test content", limit=1)
        assert len(results) > 0

    # Context manager should have automatically closed the connection
    # We can't easily test that the connection is closed without accessing internals,
    # but the test passing means the context manager methods work correctly


@pytest.mark.asyncio
async def test_client_import_document_with_custom_chunks(temp_db_path):
    """Test importing a document with pre-created chunks."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create some custom chunks with and without embeddings
        chunks = [
            Chunk(
                content="This is the first chunk",
                metadata={"custom": "metadata1"},
                order=0,
            ),
            Chunk(
                content="This is the second chunk",
                metadata={"custom": "metadata2"},
                embedding=[0.1] * Config.embeddings.model.vector_dim,
                order=1,
            ),  # With embedding
            Chunk(
                content="This is the third chunk",
                metadata={"custom": "metadata3"},
                order=2,
            ),
        ]

        # Import document with custom chunks
        document = await client.import_document(
            content="Full document content", chunks=chunks
        )

        assert document.id is not None
        assert document.content == "Full document content"

        # Verify the chunks were created correctly
        doc_chunks = await client.chunk_repository.get_by_document_id(document.id)
        assert len(doc_chunks) == 3

        # Check chunks have correct content, document_id, and order from list position
        for i, chunk in enumerate(doc_chunks):
            assert chunk.document_id == document.id
            assert chunk.content == chunks[i].content
            assert chunk.order == i  # Order should be set from list position
            assert (
                chunk.metadata["custom"] == f"metadata{i + 1}"
            )  # Original metadata preserved


@pytest.mark.asyncio
async def test_client_ask(monkeypatch, temp_db_path):
    """Test asking questions returns answer and citations."""
    from pydantic_ai.models.test import TestModel

    # Mock get_model to return TestModel
    monkeypatch.setattr(
        "haiku.rag.utils.get_model", lambda *args, **kwargs: TestModel()
    )

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a test document for the agent to search
        await client.create_document(
            content="Python is a high-level programming language.", uri="test.txt"
        )

        # Use real QA agent with TestModel
        answer, citations = await client.ask("What is Python?")

        # TestModel will generate a valid string response
        assert answer is not None
        assert isinstance(answer, str)
        assert isinstance(citations, list)


@pytest.mark.asyncio
async def test_client_expand_context(temp_db_path):
    """Test expanding search results with adjacent chunks."""
    from haiku.rag.store.models import SearchResult

    # Mock Config to have CONTEXT_CHUNK_RADIUS = 2
    with patch("haiku.rag.client.Config.processing.context_chunk_radius", 2):
        async with HaikuRAG(temp_db_path, create=True) as client:
            # Create chunks manually with precomputed embeddings to avoid network
            dim = client.chunk_repository.embedder._vector_dim
            z = [0.0] * dim
            manual_chunks = [
                Chunk(content="Chunk 0 content", order=0, embedding=z),
                Chunk(content="Chunk 1 content", order=1, embedding=z),
                Chunk(content="Chunk 2 content", order=2, embedding=z),
                Chunk(content="Chunk 3 content", order=3, embedding=z),
                Chunk(content="Chunk 4 content", order=4, embedding=z),
            ]

        doc = await client.import_document(
            content="Full document content",
            chunks=manual_chunks,
            uri="test_doc.txt",
            title="test_doc_title",
        )

        # Get all chunks for the document
        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(chunks) == 5

        # Find the middle chunk (order=2) and convert to SearchResult
        middle_chunk = next(c for c in chunks if c.order == 2)
        search_results = [SearchResult.from_chunk(middle_chunk, 0.8)]

        # Test expand_context with radius=2 and document title preserved
        expanded_results = await client.expand_context(search_results, radius=2)

        assert len(expanded_results) == 1
        expanded = expanded_results[0]

        # Check that the expanded result has combined content and preserves title/uri
        assert expanded.score == 0.8
        assert "Chunk 2 content" in expanded.content
        assert expanded.document_title == "test_doc_title"
        assert expanded.document_uri == "test_doc.txt"

        # Should include all chunks (radius=2 from chunk 2 = chunks 0,1,2,3,4)
        assert "Chunk 0 content" in expanded.content
        assert "Chunk 1 content" in expanded.content
        assert "Chunk 2 content" in expanded.content
        assert "Chunk 3 content" in expanded.content
        assert "Chunk 4 content" in expanded.content


@pytest.mark.asyncio
async def test_client_expand_context_radius_zero(temp_db_path):
    """Test expand_context with radius 0 returns original results."""
    from haiku.rag.store.models import SearchResult

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a simple document
        doc = await client.create_document(content="Simple test content")
        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)

        search_results = [SearchResult.from_chunk(chunks[0], 0.9)]
        expanded_results = await client.expand_context(search_results, radius=0)

        # Should return exactly the same results
        assert len(expanded_results) == 1
        assert expanded_results[0].content == search_results[0].content
        assert expanded_results[0].score == search_results[0].score


@pytest.mark.asyncio
async def test_client_expand_context_multiple_chunks(temp_db_path):
    """Test expand_context with multiple search results."""
    from haiku.rag.store.models import SearchResult

    with patch("haiku.rag.client.Config.processing.context_chunk_radius", 1):
        async with HaikuRAG(temp_db_path, create=True) as client:
            # Create first document with manual chunks
            doc1_chunks = [
                Chunk(content="Doc1 Part A", order=0),
                Chunk(content="Doc1 Part B", order=1),
                Chunk(content="Doc1 Part C", order=2),
            ]
            doc1 = await client.import_document(
                content="Doc1 content", chunks=doc1_chunks, uri="doc1.txt"
            )

            # Create second document with manual chunks
            doc2_chunks = [
                Chunk(content="Doc2 Section X", order=0),
                Chunk(content="Doc2 Section Y", order=1),
            ]
            doc2 = await client.import_document(
                content="Doc2 content", chunks=doc2_chunks, uri="doc2.txt"
            )

        assert doc1.id is not None
        assert doc2.id is not None
        chunks1 = await client.chunk_repository.get_by_document_id(doc1.id)
        chunks2 = await client.chunk_repository.get_by_document_id(doc2.id)

        # Get middle chunk from doc1 (order=1) and first chunk from doc2 (order=0)
        chunk1 = next(c for c in chunks1 if c.order == 1)
        chunk2 = next(c for c in chunks2 if c.order == 0)

        search_results = [
            SearchResult.from_chunk(chunk1, 0.8),
            SearchResult.from_chunk(chunk2, 0.7),
        ]
        expanded_results = await client.expand_context(search_results, radius=1)

        assert len(expanded_results) == 2

        # Check first expanded result (should include chunks 0,1,2 from doc1)
        expanded1 = expanded_results[0]
        assert expanded1.score == 0.8
        assert "Doc1 Part A" in expanded1.content
        assert "Doc1 Part B" in expanded1.content
        assert "Doc1 Part C" in expanded1.content

        # Check second expanded result (should include chunks 0,1 from doc2)
        expanded2 = expanded_results[1]
        assert expanded2.score == 0.7
        assert "Doc2 Section X" in expanded2.content
        assert "Doc2 Section Y" in expanded2.content


@pytest.mark.asyncio
async def test_client_expand_context_merges_overlapping_chunks(temp_db_path):
    """Test that overlapping expanded chunks are merged into one."""
    from haiku.rag.store.models import SearchResult

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create document with 5 chunks
        manual_chunks = [
            Chunk(content="Chunk 0", order=0),
            Chunk(content="Chunk 1", order=1),
            Chunk(content="Chunk 2", order=2),
            Chunk(content="Chunk 3", order=3),
            Chunk(content="Chunk 4", order=4),
        ]

        doc = await client.import_document(
            content="Full document content", chunks=manual_chunks
        )

        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)

        # Get adjacent chunks (orders 1 and 2) - these will overlap when expanded
        chunk1 = next(c for c in chunks if c.order == 1)
        chunk2 = next(c for c in chunks if c.order == 2)

        # With radius=1:
        # chunk1 expanded would be [0,1,2]
        # chunk2 expanded would be [1,2,3]
        # These should merge into one chunk containing [0,1,2,3]
        search_results = [
            SearchResult.from_chunk(chunk1, 0.8),
            SearchResult.from_chunk(chunk2, 0.7),
        ]
        expanded_results = await client.expand_context(search_results, radius=1)

        # Should have only 1 merged result instead of 2 overlapping ones
        assert len(expanded_results) == 1

        merged = expanded_results[0]

        # Should contain all chunks from 0 to 3
        assert "Chunk 0" in merged.content
        assert "Chunk 1" in merged.content
        assert "Chunk 2" in merged.content
        assert "Chunk 3" in merged.content
        assert "Chunk 4" not in merged.content  # Should not include chunk 4

        # Should use the higher score (0.8)
        assert merged.score == 0.8


@pytest.mark.asyncio
async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path):
    """Test that non-overlapping expanded chunks remain separate."""
    from haiku.rag.store.models import SearchResult

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create document with chunks far apart
        manual_chunks = [
            Chunk(content="Chunk 0", order=0),
            Chunk(content="Chunk 1", order=1),
            Chunk(content="Chunk 2", order=2),
            Chunk(content="Chunk 5", order=5),  # Gap here
            Chunk(content="Chunk 6", order=6),
            Chunk(content="Chunk 7", order=7),
        ]

        doc = await client.import_document(
            content="Full document content", chunks=manual_chunks
        )

        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)

        # Get chunks by index - they will have sequential orders 0,1,2,3,4,5
        # So get chunk with order=0 and chunk with order=5 (far enough apart)
        chunk0 = next(c for c in chunks if c.order == 0)  # Content: "Chunk 0"
        chunk5 = next(
            c for c in chunks if c.order == 5
        )  # Content: "Chunk 7" but now at order 5

        # chunk0 expanded: [0,1] with radius=1 (orders 0,1)
        # chunk5 expanded: [4,5] with radius=1 (orders 4,5)
        search_results = [
            SearchResult.from_chunk(chunk0, 0.8),
            SearchResult.from_chunk(chunk5, 0.7),
        ]
        expanded_results = await client.expand_context(search_results, radius=1)

        # Should have 2 separate results
        assert len(expanded_results) == 2

        # Sort by score to ensure predictable order
        expanded_results.sort(key=lambda x: x.score, reverse=True)

        chunk0_expanded = expanded_results[0]
        chunk5_expanded = expanded_results[1]

        # First chunk (order=0) expanded should contain orders [0,1]
        # Content should be "Chunk 0" + "Chunk 1"
        assert "Chunk 0" in chunk0_expanded.content
        assert "Chunk 1" in chunk0_expanded.content
        assert (
            "Chunk 5" not in chunk0_expanded.content
        )  # Should not have chunk 7 content
        assert chunk0_expanded.score == 0.8

        # Second chunk (order=5) expanded should contain orders [4,5]
        # Content should be "Chunk 6" (order 4) + "Chunk 7" (order 5)
        assert "Chunk 6" in chunk5_expanded.content
        assert "Chunk 7" in chunk5_expanded.content
        assert "Chunk 0" not in chunk5_expanded.content
        assert chunk5_expanded.score == 0.7


@pytest.mark.asyncio
async def test_client_expand_context_with_docling_merges_overlapping(temp_db_path):
    """Test that expand_context with DoclingDocument merges overlapping results."""
    from haiku.rag.store.models import SearchResult

    # Create a document with structured content that will have doc_item_refs
    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, 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_json 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 radius that should cause overlap
            expanded = await client.expand_context(search_results, radius=3)

            # 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.asyncio
async def test_client_expand_context_docling_merges_metadata(temp_db_path):
    """Test that expand_context properly merges metadata from multiple results."""
    from haiku.rag.store.models import SearchResult

    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, 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, radius=10)

            # 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]


@pytest.mark.asyncio
async def test_client_expand_context_docling_preserves_bounding_boxes(temp_db_path):
    """Test that expand_context preserves bounding boxes from DoclingDocument."""
    from haiku.rag.store.models import SearchResult

    async with HaikuRAG(temp_db_path, create=True) as client:
        doc = await client.create_document(
            content="# Test\n\nSome content here.",
            uri="test://bboxes",
        )

        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)

        if chunks:
            search_results = [SearchResult.from_chunk(chunks[0], 0.9)]
            expanded = await client.expand_context(search_results, radius=2)

            # Expanded results should exist
            assert len(expanded) == 1
            # Bounding boxes may or may not be present depending on document
            # but the field should be accessible
            _ = expanded[0].bounding_boxes


@pytest.mark.asyncio
async def test_client_create_document_stores_docling_json(temp_db_path):
    """Test that create_document stores DoclingDocument JSON."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        doc = await client.create_document(
            content="Test content for docling storage",
            uri="test://docling",
            metadata={"test": "docling_storage"},
        )

        assert doc.id is not None
        assert doc.docling_document_json is not None
        assert doc.docling_version is not None

        # Verify JSON is valid and can be parsed
        import json

        parsed = json.loads(doc.docling_document_json)
        assert "version" in parsed
        assert parsed["version"] == doc.docling_version


@pytest.mark.asyncio
async def test_client_import_document_without_docling(temp_db_path):
    """Test that import_document without docling params does not store docling JSON."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        custom_chunks = [Chunk(content="Custom chunk", order=0)]

        doc = await client.import_document(content="Test content", chunks=custom_chunks)

        assert doc.id is not None
        # When no docling params provided, they remain None
        assert doc.docling_document_json is None
        assert doc.docling_version is None


@pytest.mark.asyncio
async def test_client_import_document_validates_docling_params(temp_db_path):
    """Test that import_document validates docling parameters."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        custom_chunks = [Chunk(content="Custom chunk", order=0)]

        # Should fail if only one docling param is provided
        with pytest.raises(ValueError, match="must both be provided"):
            await client.import_document(
                content="Test content",
                chunks=custom_chunks,
                docling_document_json='{"some": "json"}',
                # Missing docling_version
            )

        with pytest.raises(ValueError, match="must both be provided"):
            await client.import_document(
                content="Test content",
                chunks=custom_chunks,
                docling_version="1.0.0",
                # Missing docling_document_json
            )

        # Should fail with invalid JSON
        with pytest.raises(ValueError, match="Invalid docling_document_json"):
            await client.import_document(
                content="Test content",
                chunks=custom_chunks,
                docling_document_json='{"invalid": "not a docling document"}',
                docling_version="1.0.0",
            )


@pytest.mark.asyncio
async def test_client_create_document_from_file_stores_docling_json(temp_db_path):
    """Test that create_document_from_source stores DoclingDocument JSON for files."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text("Test file content")

            doc = await client.create_document_from_source(temp_path)
            assert isinstance(doc, Document)

            assert doc.id is not None
            assert doc.docling_document_json is not None
            assert doc.docling_version is not None

            # Verify the stored document also has the JSON
            retrieved = await client.get_document_by_id(doc.id)
            assert retrieved is not None
            assert retrieved.docling_document_json == doc.docling_document_json
            assert retrieved.docling_version == doc.docling_version


@pytest.mark.asyncio
async def test_client_update_document_stores_docling_json(temp_db_path):
    """Test that update_document stores DoclingDocument JSON."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create initial document
        doc = await client.create_document(content="Initial content")
        assert doc.id is not None
        original_json = doc.docling_document_json

        # Update the document
        doc.content = "Updated content"
        updated_doc = await client.update_document(doc)

        assert updated_doc.docling_document_json is not None
        assert updated_doc.docling_version is not None
        # JSON should be different because content changed
        assert updated_doc.docling_document_json != original_json


@pytest.mark.asyncio
async def test_client_update_document_fields_stores_docling_json(temp_db_path):
    """Test that update_document_fields stores DoclingDocument JSON when content changes."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create initial document
        doc = await client.create_document(content="Initial content")
        assert doc.id is not None
        original_json = doc.docling_document_json

        # Update content via update_document_fields
        updated_doc = await client.update_document_fields(
            document_id=doc.id, content="New content via fields update"
        )

        assert updated_doc.docling_document_json is not None
        assert updated_doc.docling_version is not None
        # JSON should be different because content changed
        assert updated_doc.docling_document_json != original_json


@pytest.mark.asyncio
async def test_client_update_document_fields_with_custom_chunks_no_docling_json(
    temp_db_path,
):
    """Test that update_document_fields with custom chunks does not update docling JSON."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create initial document
        doc = await client.create_document(content="Initial content")
        assert doc.id is not None
        original_json = doc.docling_document_json

        # Update with custom chunks
        custom_chunks = [Chunk(content="Custom chunk", order=0)]
        updated_doc = await client.update_document_fields(
            document_id=doc.id, content="New content", chunks=custom_chunks
        )

        # Docling JSON should remain unchanged (no conversion when custom chunks provided)
        assert updated_doc.docling_document_json == original_json


@pytest.mark.asyncio
async def test_client_update_document_fields_content_docling_mutually_exclusive(
    temp_db_path,
):
    """Test that content and docling_document_json cannot both be provided."""
    from docling_core.types.doc.document import DoclingDocument

    async with HaikuRAG(temp_db_path, create=True) as client:
        doc = await client.create_document(content="Initial content")
        assert doc.id is not None

        # Create a docling document
        from docling_core.types.doc.labels import DocItemLabel

        docling_doc = DoclingDocument(name="test")
        docling_doc.add_text(label=DocItemLabel.TEXT, text="Some text")

        with pytest.raises(ValueError, match="mutually exclusive"):
            await client.update_document_fields(
                document_id=doc.id,
                content="New content",
                docling_document_json=docling_doc.model_dump_json(),
                docling_version=docling_doc.version,
            )


@pytest.mark.asyncio
async def test_client_update_document_fields_with_docling_rechunks(temp_db_path):
    """Test that providing docling_document_json without chunks triggers rechunk."""
    from docling_core.types.doc.document import DoclingDocument

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create initial document
        doc = await client.create_document(content="Initial content")
        assert doc.id is not None
        original_chunks = await client.chunk_repository.get_by_document_id(doc.id)

        # Create a new docling document with different content
        from docling_core.types.doc.labels import DocItemLabel

        docling_doc = DoclingDocument(name="updated")
        docling_doc.add_text(
            label=DocItemLabel.TEXT,
            text="Completely different text from docling document",
        )

        # Update with docling document only - should rechunk from it
        updated_doc = await client.update_document_fields(
            document_id=doc.id,
            docling_document_json=docling_doc.model_dump_json(),
            docling_version=docling_doc.version,
        )

        # Content should be extracted from docling document
        assert "Completely different text" in updated_doc.content
        assert updated_doc.docling_document_json == docling_doc.model_dump_json()
        assert updated_doc.docling_version == docling_doc.version

        # Chunks should be regenerated
        new_chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(new_chunks) > 0
        # Content should differ from original
        assert new_chunks[0].content != original_chunks[0].content


@pytest.mark.asyncio
async def test_client_update_document_fields_docling_with_chunks(temp_db_path):
    """Test that providing both docling_document_json and chunks stores both."""
    from docling_core.types.doc.document import DoclingDocument

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create initial document
        doc = await client.create_document(content="Initial content")
        assert doc.id is not None

        # Create a docling document
        from docling_core.types.doc.labels import DocItemLabel

        docling_doc = DoclingDocument(name="custom")
        docling_doc.add_text(label=DocItemLabel.TEXT, text="Text from docling")

        # Provide both docling and custom chunks
        custom_chunks = [
            Chunk(content="Custom chunk 1", order=0),
            Chunk(content="Custom chunk 2", order=1),
        ]

        updated_doc = await client.update_document_fields(
            document_id=doc.id,
            chunks=custom_chunks,
            docling_document_json=docling_doc.model_dump_json(),
            docling_version=docling_doc.version,
        )

        # Content should be extracted from docling (since content wasn't provided)
        assert "Text from docling" in updated_doc.content
        assert updated_doc.docling_document_json == docling_doc.model_dump_json()

        # Custom chunks should be used (not rechunked from docling)
        chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(chunks) == 2
        assert chunks[0].content == "Custom chunk 1"
        assert chunks[1].content == "Custom chunk 2"


@pytest.mark.asyncio
async def test_client_file_update_stores_docling_json(temp_db_path):
    """Test that updating a file re-stores DoclingDocument JSON."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.txt"
            temp_path.write_text("Original content")

            # Create initial document
            doc1 = await client.create_document_from_source(temp_path)
            assert isinstance(doc1, Document)
            original_json = doc1.docling_document_json
            original_version = doc1.docling_version

            # Modify file
            temp_path.write_text("Modified content")

            # Update document from source
            doc2 = await client.create_document_from_source(temp_path)
            assert isinstance(doc2, Document)
            assert doc2.id == doc1.id  # Same document

            # Docling JSON should be updated
            assert doc2.docling_document_json is not None
            assert doc2.docling_document_json != original_json
            assert doc2.docling_version == original_version  # Version stays same


@pytest.mark.asyncio
async def test_client_visualize_chunk_no_document(temp_db_path):
    """Test visualize_chunk returns empty list when chunk has no document_id."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        chunk = Chunk(content="Orphan chunk", order=0)
        images = await client.visualize_chunk(chunk)
        assert images == []


@pytest.mark.asyncio
async def test_client_visualize_chunk_no_docling_document(temp_db_path):
    """Test visualize_chunk returns empty list when document has no DoclingDocument."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Import document with custom chunks (no DoclingDocument)
        custom_chunks = [Chunk(content="Custom chunk", order=0)]
        doc = await client.import_document(content="Test content", chunks=custom_chunks)

        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(chunks) == 1

        images = await client.visualize_chunk(chunks[0])
        assert images == []


@pytest.mark.asyncio
async def test_client_visualize_chunk_no_bounding_boxes(temp_db_path):
    """Test visualize_chunk returns empty list when chunk has no bounding boxes."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create document from text (will have DoclingDocument but no page images)
        doc = await client.create_document(
            content="Simple text content without structure",
            uri="test://simple",
        )

        assert doc.id is not None
        assert doc.docling_document_json is not None

        chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(chunks) >= 1

        # Text documents converted via markdown won't have page images
        # so visualize_chunk should return empty list
        images = await client.visualize_chunk(chunks[0])
        assert images == []


@pytest.mark.asyncio
async def test_client_visualize_chunk_returns_list(temp_db_path):
    """Test visualize_chunk returns a list (empty or with images)."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a structured document
        markdown_content = """# Chapter 1

This is paragraph one about topic A.

This is paragraph two about topic A continued.

# Chapter 2

This is paragraph four about topic C.
"""
        doc = await client.create_document(
            content=markdown_content,
            uri="test://structured",
        )

        assert doc.id is not None
        chunks = await client.chunk_repository.get_by_document_id(doc.id)

        # Find a chunk with doc_item_refs
        chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]

        if chunks_with_refs:
            # visualize_chunk should return a list (possibly empty if no page images)
            images = await client.visualize_chunk(chunks_with_refs[0])
            assert isinstance(images, list)


@pytest.mark.integration
@pytest.mark.asyncio
async def test_client_visualize_chunk_with_pdf(temp_db_path):
    """Test visualize_chunk returns images with bounding boxes for PDF documents."""
    from PIL.Image import Image as PILImage

    pdf_path = Path("tests/data/doclaynet.pdf")
    if not pdf_path.exists():
        pytest.skip("doclaynet.pdf not found")

    async with HaikuRAG(temp_db_path, create=True) as client:
        doc = await client.create_document_from_source(pdf_path)
        assert isinstance(doc, Document)
        assert doc.id is not None
        assert doc.docling_document_json is not None

        chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(chunks) > 0

        # Find a chunk with doc_item_refs (bounding box info)
        chunks_with_refs = [c for c in chunks if c.get_chunk_metadata().doc_item_refs]
        assert len(chunks_with_refs) > 0, "PDF should have chunks with doc_item_refs"

        # Visualize a chunk - should return images with bounding boxes drawn
        images = await client.visualize_chunk(chunks_with_refs[0])

        assert isinstance(images, list)
        assert len(images) > 0, "PDF with page images should return visualizations"

        # Verify returned objects are PIL Images
        for img in images:
            assert isinstance(img, PILImage)