"""A3: SearchResult.image_data, expand_context preservation, multimodal ToolReturn."""

import base64
from dataclasses import dataclass
from unittest.mock import AsyncMock

import pytest
from pydantic_ai import RunContext
from pydantic_ai.messages import BinaryContent, ToolReturn
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage

from haiku.rag.client import HaikuRAG
from haiku.rag.client.search import _populate_image_data
from haiku.rag.config import Config
from haiku.rag.context import expand_with_items
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
from haiku.rag.tools.search import create_search_toolset

PICTURE_BYTES = b"\x89PNG\r\n\x1a\nfake-picture-bytes"
PICTURE_B64 = base64.b64encode(PICTURE_BYTES).decode("ascii")


@pytest.mark.asyncio
async def test_populate_image_data_attaches_base64(temp_db_path):
    async with HaikuRAG(temp_db_path, create=True) as rag:
        await rag.document_item_repository.create_items(
            "doc-1",
            [
                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=PICTURE_BYTES,
                ),
            ],
        )

        text_only = SearchResult(
            content="Some text",
            score=1.0,
            document_id="doc-1",
            doc_item_refs=["#/texts/0"],
            labels=["paragraph"],
        )
        with_picture = SearchResult(
            content="A figure caption",
            score=0.9,
            document_id="doc-1",
            doc_item_refs=["#/texts/0", "#/pictures/0"],
            labels=["paragraph", "picture"],
        )

        await _populate_image_data(rag, [text_only, with_picture])

        # Text-only result is unchanged
        assert text_only.image_data is None
        # Picture-bearing result has the bytes attached
        assert with_picture.image_data == {"#/pictures/0": PICTURE_B64}


@pytest.mark.asyncio
async def test_client_search_include_images_false_skips_lookup(temp_db_path):
    """include_images=False must short-circuit the picture-bytes lookup."""
    async with HaikuRAG(temp_db_path, create=True) as rag:
        await rag.document_item_repository.create_items(
            "doc-1",
            [
                DocumentItem(
                    document_id="doc-1",
                    position=0,
                    self_ref="#/pictures/0",
                    label="picture",
                    picture_data=PICTURE_BYTES,
                ),
            ],
        )
        # Spy that we never reach the picture-bytes accessor
        rag.document_item_repository.get_pictures_for_chunk = AsyncMock(  # type: ignore[method-assign]
            wraps=rag.document_item_repository.get_pictures_for_chunk
        )

        from haiku.rag.client.search import search

        # Stub the chunk-search results so we don't depend on embeddings/FTS
        async def fake_chunk_search(*args, **kwargs):
            return []

        rag.chunk_repository.search = fake_chunk_search  # type: ignore[method-assign]

        await search(rag, "anything", include_images=False)
        rag.document_item_repository.get_pictures_for_chunk.assert_not_called()


@pytest.mark.asyncio
async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_path):
    """A picture item with empty text must keep its self_ref through expansion."""
    async with HaikuRAG(temp_db_path, create=True) as rag:
        # Build an items table with a section header + a paragraph match + an
        # adjacent picture row that has no text. expand_with_items used to
        # filter the picture out via the `if item.text:` guard; A3 keeps it.
        await rag.document_item_repository.create_items(
            "doc-1",
            [
                DocumentItem(
                    document_id="doc-1",
                    position=0,
                    self_ref="#/texts/0",
                    label="section_header",
                    text="Methods",
                ),
                DocumentItem(
                    document_id="doc-1",
                    position=1,
                    self_ref="#/texts/1",
                    label="paragraph",
                    text="The figure below shows the architecture.",
                ),
                DocumentItem(
                    document_id="doc-1",
                    position=2,
                    self_ref="#/pictures/0",
                    label="picture",
                    text="",
                    picture_data=PICTURE_BYTES,
                ),
                DocumentItem(
                    document_id="doc-1",
                    position=3,
                    self_ref="#/texts/2",
                    label="paragraph",
                    text="More commentary follows.",
                ),
            ],
        )

        # Match on the paragraph that mentions the figure.
        seed = SearchResult(
            content="The figure below shows the architecture.",
            score=1.0,
            document_id="doc-1",
            doc_item_refs=["#/texts/1"],
            labels=["paragraph"],
        )
        expanded = await expand_with_items(
            rag.document_item_repository,
            "doc-1",
            [seed],
            max_chars=10_000,
        )
        assert len(expanded) == 1
        out = expanded[0]
        assert "#/pictures/0" in out.doc_item_refs, (
            "Picture self_ref should survive expansion even with empty text"
        )
        assert "picture" in out.labels


@pytest.mark.asyncio
async def test_rechunk_preserves_picture_data(temp_db_path):
    """``rebuild --rechunk`` keeps ``picture_data`` for every picture row."""
    from haiku.rag.client import RebuildMode
    from haiku.rag.client.documents import _store_document_with_chunks
    from haiku.rag.store.models.document import Document
    from tests.store.test_document_items import _docling_doc_with_picture

    docling_doc = _docling_doc_with_picture()

    async with HaikuRAG(temp_db_path, create=True) as rag:
        rag._config.processing.pictures = "image"
        document = Document(content="x", uri="test://doc")
        document.set_docling(docling_doc)
        created = await _store_document_with_chunks(rag, document, [], docling_doc)
        assert created.id is not None
        before = await rag.document_item_repository.get_all_picture_data(created.id)
        assert before.get("#/pictures/0") is not None

        async for _ in rag.rebuild_database(mode=RebuildMode.RECHUNK):
            pass

        after = await rag.document_item_repository.get_all_picture_data(created.id)
        assert after.get("#/pictures/0") == before.get("#/pictures/0")


@pytest.mark.asyncio
async def test_update_clears_picture_data_when_mode_none(temp_db_path):
    """Switching to ``pictures="none"`` and re-running update_document
    drops picture_data — the snapshot/merge gate is what gives users a
    "rebuild reclaims storage" path when they downgrade modes."""
    from haiku.rag.client.documents import (
        _store_document_with_chunks,
        _update_document_with_chunks,
    )
    from haiku.rag.store.models.document import Document
    from tests.store.test_document_items import _docling_doc_with_picture

    docling_doc = _docling_doc_with_picture()

    async with HaikuRAG(temp_db_path, create=True) as rag:
        # Ingest under "image" so picture bytes land in document_items.
        rag._config.processing.pictures = "image"
        document = Document(content="x", uri="test://doc")
        document.set_docling(docling_doc)
        created = await _store_document_with_chunks(rag, document, [], docling_doc)
        assert created.id is not None
        before = await rag.document_item_repository.get_all_picture_data(created.id)
        assert before.get("#/pictures/0") is not None

        # Downgrade to "none" and re-run update with the (already stripped)
        # docling pulled from storage. The snapshot/merge must be skipped so
        # picture_data is cleared on the new items rows.
        rag._config.processing.pictures = "none"
        from_blob = created.get_docling_document()
        assert from_blob is not None
        await _update_document_with_chunks(rag, created, [], from_blob)

        after = await rag.document_item_repository.get_all_picture_data(created.id)
        assert after.get("#/pictures/0") is None


@pytest.mark.asyncio
async def test_expand_context_repopulates_image_data(temp_db_path):
    """expand_context rebuilds SearchResult objects via expand_with_items, so
    it must re-attach picture bytes — otherwise vision flows downstream see
    empty image_data after expansion."""
    async with HaikuRAG(temp_db_path, create=True) as rag:
        await rag.document_item_repository.create_items(
            "doc-1",
            [
                DocumentItem(
                    document_id="doc-1",
                    position=0,
                    self_ref="#/texts/0",
                    label="section_header",
                    text="Methods",
                ),
                DocumentItem(
                    document_id="doc-1",
                    position=1,
                    self_ref="#/texts/1",
                    label="paragraph",
                    text="The figure below shows the architecture.",
                ),
                DocumentItem(
                    document_id="doc-1",
                    position=2,
                    self_ref="#/pictures/0",
                    label="picture",
                    text="",
                    picture_data=PICTURE_BYTES,
                ),
            ],
        )
        seed = SearchResult(
            content="The figure below shows the architecture.",
            score=1.0,
            document_id="doc-1",
            doc_item_refs=["#/texts/1"],
            labels=["paragraph"],
            image_data=None,
        )
        expanded = await rag.expand_context([seed])
        assert len(expanded) == 1
        out = expanded[0]
        assert "#/pictures/0" in out.doc_item_refs
        assert out.image_data == {"#/pictures/0": PICTURE_B64}


@dataclass
class _Deps:
    client: object


@pytest.mark.asyncio
async def test_search_tool_returns_multimodal_when_picture_present():
    """The agent-facing search tool must wrap text + BinaryContent in ToolReturn
    whenever a result carries picture image_data."""
    picture_result = SearchResult(
        content="A diagram of the layout",
        score=1.0,
        chunk_id="chunk-1",
        document_id="doc-1",
        doc_item_refs=["#/pictures/0"],
        labels=["picture"],
        image_data={"#/pictures/0": PICTURE_B64},
    )

    fake_client = AsyncMock()
    fake_client.search = AsyncMock(return_value=[picture_result])
    fake_client.expand_context = AsyncMock(return_value=[picture_result])

    toolset = create_search_toolset(Config, expand_context=False)
    func = toolset.tools["search"].function

    ctx = RunContext(
        deps=_Deps(client=fake_client),  # type: ignore[arg-type]
        model=TestModel(),
        usage=RunUsage(),
        run_id="run-1",
    )
    result = await func(ctx, "anything")

    assert isinstance(result, ToolReturn)
    assert isinstance(result.return_value, str)
    assert "Type: picture" in result.return_value or "rank 1" in result.return_value
    assert result.content is not None
    assert len(result.content) == 1
    part = result.content[0]
    assert isinstance(part, BinaryContent)
    assert part.media_type == "image/png"
    assert part.identifier == "#/pictures/0"
    assert part.data == PICTURE_BYTES


@pytest.mark.asyncio
async def test_search_tool_returns_plain_string_when_no_pictures():
    """When no result carries image_data the tool returns a plain str (no
    ToolReturn wrapper) so non-vision flows are unaffected."""
    text_result = SearchResult(
        content="Some text",
        score=1.0,
        chunk_id="chunk-1",
        document_id="doc-1",
        doc_item_refs=["#/texts/0"],
        labels=["paragraph"],
    )

    fake_client = AsyncMock()
    fake_client.search = AsyncMock(return_value=[text_result])
    fake_client.expand_context = AsyncMock(return_value=[text_result])

    toolset = create_search_toolset(Config, expand_context=False)
    func = toolset.tools["search"].function

    ctx = RunContext(
        deps=_Deps(client=fake_client),  # type: ignore[arg-type]
        model=TestModel(),
        usage=RunUsage(),
        run_id="run-1",
    )
    result = await func(ctx, "anything")

    assert isinstance(result, str)
    assert "rank 1" in result