import pytest

from haiku.rag.client import HaikuRAG
from haiku.rag.client.documents import (
    _store_document_with_chunks,
    _update_document_with_chunks,
)
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document_item import (
    DocumentItem,
    extract_item_text,
    extract_items,
)
from haiku.rag.store.repositories.document_item import DocumentItemRepository


def _make_docling_doc():
    """Create a DoclingDocument with mixed item types for testing."""
    from docling_core.types.doc.document import DoclingDocument, TableData
    from docling_core.types.doc.labels import DocItemLabel

    doc = DoclingDocument(name="test")
    doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Introduction")
    doc.add_text(label=DocItemLabel.PARAGRAPH, text="This is the first paragraph.")
    doc.add_text(label=DocItemLabel.PARAGRAPH, text="This is the second paragraph.")
    doc.add_table(data=TableData(num_rows=2, num_cols=2, table_cells=[]))
    doc.add_text(label=DocItemLabel.SECTION_HEADER, text="Conclusion")
    doc.add_text(label=DocItemLabel.PARAGRAPH, text="Final thoughts here.")
    return doc


class TestExtractItems:
    def test_extracts_all_items(self):
        doc = _make_docling_doc()
        items = extract_items("doc-1", doc)

        assert len(items) == 6
        assert all(item.document_id == "doc-1" for item in items)
        assert [item.position for item in items] == [0, 1, 2, 3, 4, 5]

    def test_extracts_labels(self):
        doc = _make_docling_doc()
        items = extract_items("doc-1", doc)

        assert items[0].label == "section_header"
        assert items[1].label == "paragraph"
        assert items[3].label == "table"
        assert items[4].label == "section_header"

    def test_extracts_text(self):
        doc = _make_docling_doc()
        items = extract_items("doc-1", doc)

        assert items[0].text == "Introduction"
        assert items[1].text == "This is the first paragraph."
        assert items[5].text == "Final thoughts here."

    def test_extracts_self_refs(self):
        doc = _make_docling_doc()
        items = extract_items("doc-1", doc)

        assert all(item.self_ref.startswith("#/") for item in items)

    def test_table_gets_markdown_text(self):
        doc = _make_docling_doc()
        items = extract_items("doc-1", doc)

        table_item = items[3]
        assert table_item.label == "table"
        # Table should have some text from export_to_markdown
        assert isinstance(table_item.text, str)


class TestExtractItemText:
    def test_text_item(self):
        from docling_core.types.doc.document import DoclingDocument
        from docling_core.types.doc.labels import DocItemLabel

        doc = DoclingDocument(name="test")
        doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
        item, _ = next(iter(doc.iterate_items()))
        assert extract_item_text(item, doc) == "Hello world"

    def test_returns_none_for_empty_item(self):
        from docling_core.types.doc.document import DoclingDocument

        doc = DoclingDocument(name="test")
        # An empty doc has no items to extract text from
        items = extract_items("doc-1", doc)
        assert items == []


@pytest.mark.asyncio
class TestDocumentItemRepository:
    async def test_create_and_get_range(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            repo = DocumentItemRepository(rag.store)

            items = [
                DocumentItem(
                    document_id="doc-1",
                    position=i,
                    self_ref=f"#/texts/{i}",
                    label="paragraph",
                    text=f"Item {i}",
                    page_numbers=[1],
                )
                for i in range(10)
            ]
            await repo.create_items("doc-1", items)

            result = await repo.get_items_in_range("doc-1", 3, 7)
            assert len(result) == 5
            assert result[0].position == 3
            assert result[-1].position == 7
            assert result[0].text == "Item 3"

    async def test_resolve_refs(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            repo = DocumentItemRepository(rag.store)

            items = [
                DocumentItem(
                    document_id="doc-1",
                    position=i,
                    self_ref=f"#/texts/{i}",
                    label="paragraph",
                    text=f"Item {i}",
                )
                for i in range(10)
            ]
            await repo.create_items("doc-1", items)

            refs = await repo.resolve_refs(
                "doc-1", ["#/texts/2", "#/texts/7", "#/texts/999"]
            )
            assert refs == {"#/texts/2": 2, "#/texts/7": 7}

    async def test_get_item_count(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            repo = DocumentItemRepository(rag.store)

            items = [
                DocumentItem(
                    document_id="doc-1",
                    position=i,
                    self_ref=f"#/texts/{i}",
                    label="paragraph",
                    text=f"Item {i}",
                )
                for i in range(15)
            ]
            await repo.create_items("doc-1", items)

            assert await repo.get_item_count("doc-1") == 15
            assert await repo.get_item_count("nonexistent") == 0

    async def test_delete_by_document_id(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            repo = DocumentItemRepository(rag.store)

            for doc_id in ["doc-1", "doc-2"]:
                items = [
                    DocumentItem(
                        document_id=doc_id,
                        position=i,
                        self_ref=f"#/texts/{i}",
                        label="paragraph",
                        text=f"Item {i}",
                    )
                    for i in range(5)
                ]
                await repo.create_items(doc_id, items)

            await repo.delete_by_document_id("doc-1")

            assert await repo.get_item_count("doc-1") == 0
            assert await repo.get_item_count("doc-2") == 5

    async def test_empty_refs_returns_empty(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            repo = DocumentItemRepository(rag.store)
            assert await repo.resolve_refs("doc-1", []) == {}

    async def test_items_sorted_by_position(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as rag:
            repo = DocumentItemRepository(rag.store)

            # Insert in reverse order
            items = [
                DocumentItem(
                    document_id="doc-1",
                    position=9 - i,
                    self_ref=f"#/texts/{9 - i}",
                    label="paragraph",
                    text=f"Item {9 - i}",
                )
                for i in range(10)
            ]
            await repo.create_items("doc-1", items)

            result = await repo.get_items_in_range("doc-1", 0, 9)
            positions = [item.position for item in result]
            assert positions == sorted(positions)


@pytest.mark.asyncio
class TestDocumentItemPopulation:
    async def test_store_document_populates_items(self, temp_db_path):
        """Test that _store_document_with_chunks populates items when given a docling_document."""
        from haiku.rag.store.models.document import Document

        docling_doc = _make_docling_doc()

        async with HaikuRAG(temp_db_path, create=True) as rag:
            document = Document(
                content="test content",
                uri="test://doc",
            )
            document.set_docling(docling_doc)

            # Use _store_document_with_chunks directly with empty chunks
            # to avoid needing embeddings
            created = await _store_document_with_chunks(rag, document, [], docling_doc)
            assert created.id is not None

            count = await rag.document_item_repository.get_item_count(created.id)
            assert count == 6

            items = await rag.document_item_repository.get_items_in_range(
                created.id, 0, count
            )
            assert items[0].label == "section_header"
            assert items[0].text == "Introduction"
            assert items[1].label == "paragraph"

    async def test_update_document_replaces_items(self, temp_db_path):
        """Test that _update_document_with_chunks replaces items."""
        from docling_core.types.doc.document import DoclingDocument
        from docling_core.types.doc.labels import DocItemLabel

        from haiku.rag.store.models.document import Document

        docling_doc = _make_docling_doc()

        async with HaikuRAG(temp_db_path, create=True) as rag:
            document = Document(
                content="test content",
                uri="test://doc",
            )
            document.set_docling(docling_doc)
            created = await _store_document_with_chunks(rag, document, [], docling_doc)
            assert created.id is not None
            assert await rag.document_item_repository.get_item_count(created.id) == 6

            # Update with a simpler document
            new_doc = DoclingDocument(name="updated")
            new_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Only one item now.")
            created.set_docling(new_doc)

            await _update_document_with_chunks(rag, created, [], new_doc)
            assert await rag.document_item_repository.get_item_count(created.id) == 1

    async def test_delete_document_cascades_items(self, temp_db_path):
        """Test that deleting a document also deletes its items."""
        from haiku.rag.store.models.document import Document

        docling_doc = _make_docling_doc()

        async with HaikuRAG(temp_db_path, create=True) as rag:
            document = Document(
                content="test content",
                uri="test://doc",
            )
            document.set_docling(docling_doc)
            created = await _store_document_with_chunks(rag, document, [], docling_doc)
            assert created.id is not None
            assert await rag.document_item_repository.get_item_count(created.id) == 6

            await rag.delete_document(created.id)
            assert await rag.document_item_repository.get_item_count(created.id) == 0


@pytest.mark.asyncio
class TestDocumentItemMigration:
    async def test_migration_populates_items_for_existing_documents(self, temp_db_path):
        """Test that the v0.40.0 migration populates items for pre-existing documents."""
        from haiku.rag.store.compression import compress_docling_split
        from haiku.rag.store.engine import DocumentRecord

        docling_doc = _make_docling_doc()
        json_str = docling_doc.model_dump_json()
        structure, pages = compress_docling_split(json_str)

        # Create a database at a pre-migration version with a document
        async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
            await store.set_haiku_version("0.39.0")
            doc_record = DocumentRecord(
                id="test-doc-1",
                content="test content",
                uri="test://doc",
                docling_document=structure,
                docling_pages=pages,
                docling_version=docling_doc.version,
            )
            await store.documents_table.add([doc_record])

            # Verify no items exist yet
            assert await store.document_items_table.count_rows() == 0

        # Re-open with skip_migration_check and run migration
        async with Store(temp_db_path, skip_migration_check=True) as store:
            applied = await store.migrate()

            # Should have applied the v0.40.0 migration
            assert any("document_items" in desc for desc in applied)

            # Items should now exist
            item_count = await store.document_items_table.count_rows(
                filter="document_id = 'test-doc-1'"
            )
            assert item_count == 6

            # Verify item content
            items = await (
                store.document_items_table.query()
                .where("document_id = 'test-doc-1'")
                .to_list()
            )
            labels = {row["label"] for row in items}
            assert "section_header" in labels
            assert "paragraph" in labels
            assert "table" in labels

    async def test_migration_skips_documents_without_docling(self, temp_db_path):
        """Test that migration handles documents without docling data."""
        from haiku.rag.store.engine import DocumentRecord

        async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
            await store.set_haiku_version("0.39.0")
            doc_record = DocumentRecord(
                id="no-docling",
                content="plain text document",
            )
            await store.documents_table.add([doc_record])

        async with Store(temp_db_path, skip_migration_check=True) as store:
            await store.migrate()

            # No items should have been created
            assert await store.document_items_table.count_rows() == 0


@pytest.mark.asyncio
class TestPictureDataStorage:
    async def test_create_and_get_picture_bytes(self, temp_db_path):
        """Round-trip picture bytes through DocumentItem and the repository."""
        async with HaikuRAG(temp_db_path, create=True) as rag:
            repo = DocumentItemRepository(rag.store)

            png_bytes = b"\x89PNG\r\n\x1a\nfake-picture-bytes"
            items = [
                DocumentItem(
                    document_id="doc-1",
                    position=0,
                    self_ref="#/texts/0",
                    label="paragraph",
                    text="Some text",
                ),
                DocumentItem(
                    document_id="doc-1",
                    position=1,
                    self_ref="#/pictures/0",
                    label="picture",
                    text="",
                    picture_data=png_bytes,
                ),
            ]
            await repo.create_items("doc-1", items)

            # Single-ref lookup
            got = await repo.get_picture_bytes("doc-1", "#/pictures/0")
            assert got == png_bytes
            # Missing ref returns None
            assert await repo.get_picture_bytes("doc-1", "#/pictures/999") is None
            # Non-picture row has no bytes
            assert await repo.get_picture_bytes("doc-1", "#/texts/0") is None

            # Batch lookup omits refs without bytes
            batch = await repo.get_pictures_for_chunk(
                "doc-1", ["#/pictures/0", "#/texts/0", "#/pictures/999"]
            )
            assert batch == {"#/pictures/0": png_bytes}

            # Empty refs returns empty dict
            assert await repo.get_pictures_for_chunk("doc-1", []) == {}

    async def test_hot_paths_exclude_picture_data(self, temp_db_path):
        """Light read paths must NOT pull picture_data into memory."""
        async with HaikuRAG(temp_db_path, create=True) as rag:
            repo = DocumentItemRepository(rag.store)

            heavy = b"x" * 1024
            await repo.create_items(
                "doc-1",
                [
                    DocumentItem(
                        document_id="doc-1",
                        position=0,
                        self_ref="#/pictures/0",
                        label="picture",
                        text="",
                        picture_data=heavy,
                    ),
                ],
            )

            for item in await repo.get_all_items("doc-1"):
                assert item.picture_data is None
            for item in await repo.get_items_in_range("doc-1", 0, 10):
                assert item.picture_data is None
            grouped = await repo.get_all_items_grouped(["doc-1"])
            for item in grouped["doc-1"]:
                assert item.picture_data is None

            # But the picture-byte accessors still work
            assert (await repo.get_picture_bytes("doc-1", "#/pictures/0")) == heavy

    async def test_fresh_db_has_picture_data_column(self, temp_db_path):
        """A newly-created DB has picture_data on document_items via _init_tables."""
        async with HaikuRAG(temp_db_path, create=True) as rag:
            schema = await rag.store.document_items_table.schema()
            assert "picture_data" in {f.name for f in schema}