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
        # Content is stored as markdown export, check key text is preserved
        assert "Jakarta" in created_doc.content
        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 "Jakarta" in retrieved_doc.content
        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 "Jakarta" in retrieved_by_uri.content

        # 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
        updated_doc = await client.update_document(
            document_id=retrieved_doc.id,  # type: ignore[arg-type]
            content="Updated content",
        )
        assert updated_doc.content == "Updated content"

        # 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(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(
            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(
            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(
            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(
            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 that expand_context method exists and works with basic input."""
    from haiku.rag.store.models import SearchResult

    async with HaikuRAG(temp_db_path, create=True) as client:
        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)

        assert len(expanded_results) == 1
        assert expanded_results[0].score == 0.9


@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",
            )

        # Should fail if neither content nor docling_document_json provided
        with pytest.raises(ValueError, match="Either content or docling_document_json"):
            await client.import_document(chunks=custom_chunks)


@pytest.mark.asyncio
async def test_client_import_document_extracts_content_from_docling(temp_db_path):
    """Test that import_document extracts content from DoclingDocument when not provided."""
    from docling_core.types.doc.document import DoclingDocument
    from docling_core.types.doc.labels import DocItemLabel

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create a docling document with some content
        docling_doc = DoclingDocument(name="test")
        docling_doc.add_text(
            label=DocItemLabel.TEXT, text="Content from docling document"
        )

        custom_chunks = [Chunk(content="Chunk content", order=0)]

        # Import without content - should extract from docling
        doc = await client.import_document(
            chunks=custom_chunks,
            docling_document_json=docling_doc.model_dump_json(),
            docling_version=docling_doc.version,
        )

        assert doc.id is not None
        assert "Content from docling document" in doc.content
        assert doc.docling_document_json == docling_doc.model_dump_json()


@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 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
        updated_doc = await client.update_document(
            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_with_custom_chunks_no_docling_json(
    temp_db_path,
):
    """Test that update_document 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(
            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_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(
                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_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(
            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_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(
            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)


# =============================================================================
# convert() method tests
# =============================================================================


@pytest.mark.asyncio
async def test_client_convert_text(temp_db_path):
    """Test convert() with plain text content."""
    from docling_core.types.doc.document import DoclingDocument

    async with HaikuRAG(temp_db_path, create=True) as client:
        text = "This is some test content for conversion."
        docling_doc = await client.convert(text)

        assert isinstance(docling_doc, DoclingDocument)
        # Check the content is preserved in markdown export
        markdown = docling_doc.export_to_markdown()
        assert "test content" in markdown


@pytest.mark.asyncio
async def test_client_convert_file(temp_db_path):
    """Test convert() with a file path."""
    from docling_core.types.doc.document import DoclingDocument

    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("File content for conversion test.")

            docling_doc = await client.convert(temp_path)

            assert isinstance(docling_doc, DoclingDocument)
            markdown = docling_doc.export_to_markdown()
            assert "File content" in markdown


@pytest.mark.asyncio
async def test_client_convert_file_not_found(temp_db_path):
    """Test convert() raises ValueError for non-existent file."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with pytest.raises(ValueError, match="File does not exist"):
            await client.convert(Path("/nonexistent/path/file.txt"))


@pytest.mark.asyncio
async def test_client_convert_unsupported_extension(temp_db_path):
    """Test convert() raises ValueError for unsupported file extension."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        with tempfile.TemporaryDirectory() as temp_dir:
            temp_path = Path(temp_dir) / "test.xyz"
            temp_path.write_text("content")

            with pytest.raises(ValueError, match="Unsupported file extension"):
                await client.convert(temp_path)


@pytest.mark.asyncio
async def test_client_convert_file_uri(temp_db_path):
    """Test convert() with a file:// URI string."""
    from docling_core.types.doc.document import DoclingDocument

    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("URI file content.")
            file_uri = temp_path.as_uri()

            docling_doc = await client.convert(file_uri)

            assert isinstance(docling_doc, DoclingDocument)
            markdown = docling_doc.export_to_markdown()
            assert "URI file content" in markdown


# =============================================================================
# chunk() method tests
# =============================================================================


@pytest.mark.asyncio
async def test_client_chunk_basic(temp_db_path):
    """Test chunk() produces Chunk objects from DoclingDocument."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # First convert some text
        docling_doc = await client.convert("This is test content for chunking.")

        # Then chunk it
        chunks = await client.chunk(docling_doc)

        assert isinstance(chunks, list)
        assert len(chunks) > 0
        assert all(isinstance(c, Chunk) for c in chunks)
        # Chunks should have content but no embedding yet
        assert all(c.content for c in chunks)
        assert all(c.embedding is None for c in chunks)
        # Chunks should not have document_id yet (not stored)
        assert all(c.document_id is None for c in chunks)


@pytest.mark.asyncio
async def test_client_chunk_preserves_metadata(temp_db_path):
    """Test chunk() preserves structured metadata from DoclingDocument."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Convert structured markdown
        markdown = """# Chapter 1

This is the first paragraph.

## Section 1.1

This is a subsection.
"""
        docling_doc = await client.convert(markdown)
        chunks = await client.chunk(docling_doc)

        assert len(chunks) > 0

        # Check that at least some chunks have metadata
        has_metadata = False
        for chunk in chunks:
            meta = chunk.get_chunk_metadata()
            if meta.doc_item_refs or meta.headings:
                has_metadata = True
                break

        assert has_metadata, "Chunks should have structured metadata"


@pytest.mark.asyncio
async def test_client_chunk_empty_document(temp_db_path):
    """Test chunk() with empty DoclingDocument."""
    from docling_core.types.doc.document import DoclingDocument

    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create an empty DoclingDocument
        empty_doc = DoclingDocument(name="empty")

        chunks = await client.chunk(empty_doc)

        assert isinstance(chunks, list)
        assert len(chunks) == 0


@pytest.mark.asyncio
async def test_import_document_embeds_chunks_without_embeddings(temp_db_path):
    """Test that import_document embeds chunks that don't have embeddings."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        # Create chunks without embeddings
        chunks = [
            Chunk(content="First chunk without embedding", order=0),
            Chunk(content="Second chunk without embedding", order=1),
        ]

        # Import document with chunks that have no embeddings
        doc = await client.import_document(
            content="Document with unembedded chunks",
            chunks=chunks,
        )
        assert doc.id is not None

        # Verify chunks were stored
        stored_chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(stored_chunks) == 2

        # Verify vector search works (proves embeddings were generated)
        results = await client.search("First chunk", search_type="vector")
        assert len(results) > 0
        assert results[0].content == "First chunk without embedding"


@pytest.mark.asyncio
async def test_update_document_embeds_chunks_without_embeddings(temp_db_path):
    """Test that update_document embeds chunks that don't have embeddings."""
    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

        # Update with chunks that have no embeddings
        new_chunks = [
            Chunk(content="Updated chunk without embedding", order=0),
        ]
        await client.update_document(
            document_id=doc.id,
            content="Updated content",
            chunks=new_chunks,
        )

        # Verify chunks were stored
        stored_chunks = await client.chunk_repository.get_by_document_id(doc.id)
        assert len(stored_chunks) == 1
        assert stored_chunks[0].content == "Updated chunk without embedding"

        # Verify vector search works (proves embeddings were generated)
        results = await client.search("Updated chunk", search_type="vector")
        assert len(results) > 0
        assert results[0].content == "Updated chunk without embedding"


@pytest.mark.asyncio
async def test_client_create_document_with_html_format(temp_db_path):
    """Test create_document with HTML format preserves document structure."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        html_content = """
        <h1>Main Title</h1>
        <p>Introduction paragraph.</p>
        <h2>Section Header</h2>
        <ul>
            <li>Item 1</li>
            <li>Item 2</li>
        </ul>
        """

        doc = await client.create_document(
            content=html_content,
            uri="test://html-doc",
            format="html",
        )

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

        # Verify the DoclingDocument has proper structure
        docling_doc = doc.get_docling_document()
        assert docling_doc is not None

        items = list(docling_doc.iterate_items())
        labels = [str(getattr(item, "label", "")) for item, _ in items]

        # HTML format should preserve headers and list items
        assert "title" in labels or "section_header" in labels
        assert "list_item" in labels


@pytest.mark.asyncio
async def test_client_convert_with_html_format(temp_db_path):
    """Test convert with HTML format."""
    async with HaikuRAG(temp_db_path, create=True) as client:
        html_content = "<h1>Title</h1><p>Text</p>"

        docling_doc = await client.convert(html_content, format="html")

        items = list(docling_doc.iterate_items())
        labels = [str(getattr(item, "label", "")) for item, _ in items]

        assert "title" in labels