from __future__ import annotations

import pytest
from pydantic import ValidationError

from deerflow.config.llmwiki_config import LocalWikiIndexConfig
from deerflow.integrations.weknora.local_index.normalize import normalize_wiki_page
from deerflow.integrations.weknora.local_index.search import WikiIndexError, WikiVectorSearchService
from deerflow.integrations.weknora.local_index.vector_cache import WikiVectorCache
from deerflow.integrations.weknora.local_index.vector_codec import encode_vector
from deerflow.persistence.llmwiki_index.memory import MemoryLlmWikiIndexStore


class FakeEmbedding:
    dimensions = 2
    fingerprint = "fp"

    async def embed_query(self, text: str) -> list[float]:
        return [1.0, 0.0]


def test_enabled_index_requires_supported_enabled_embedding() -> None:
    with pytest.raises(ValidationError, match="embedding must be enabled"):
        LocalWikiIndexConfig(enabled=True, embedding={"enabled": False})
    with pytest.raises(ValidationError, match="provider must be openai"):
        LocalWikiIndexConfig(embedding={"provider": "custom"})


@pytest.mark.asyncio
async def test_exact_search_groups_sections_by_wiki_page_and_filters_drafts() -> None:
    store = MemoryLlmWikiIndexStore()
    await store.try_acquire_sync_lease("kb", owner="w", sync_id="s", lease_seconds=60, fingerprint="fp")
    published = normalize_wiki_page("kb", {"slug": "alice", "title": "Alice", "summary": "Founder", "content": "Alpha", "status": "published"})
    draft = normalize_wiki_page("kb", {"slug": "secret", "title": "Secret", "content": "Draft", "status": "draft"})
    a, _ = encode_vector([1.0, 0.0], dimensions=2)
    b, _ = encode_vector([0.9, 0.1], dimensions=2)
    await store.replace_page_vectors(
        page=published,
        vectors=[
            {"section_index": 0, "heading": "A", "section_content": "Alpha", "content_hash": "a", "embedding_dimensions": 2, "vector_blob": a},
            {"section_index": 1, "heading": "B", "section_content": "Beta", "content_hash": "b", "embedding_dimensions": 2, "vector_blob": b},
        ],
        fingerprint="fp",
        sync_id="s",
    )
    await store.replace_page_vectors(page=draft, vectors=[{"section_index": 0, "heading": None, "section_content": "Draft", "content_hash": "d", "embedding_dimensions": 2, "vector_blob": a}], fingerprint="fp", sync_id="s")
    config = LocalWikiIndexConfig(min_similarity=0, top_k_sections=10, top_k_pages=8, max_sections_per_page=2)
    service = WikiVectorSearchService(config, FakeEmbedding(), WikiVectorCache(store, revision_check_seconds=0))
    mapping = {"id": "kb", "name": "People"}

    result = await service.search("alice", [mapping])
    assert [item["wiki_slug"] for item in result["results"]] == ["alice"]
    assert len(result["results"][0]["matched_sections"]) == 2
    assert "wiki_slug=alice" in result["results"][0]["url"]

    with_drafts = await service.search("alice", [mapping], include_drafts=True)
    assert {item["wiki_slug"] for item in with_drafts["results"]} == {"alice", "secret"}

    section_result = await service.search_sections("alice", [mapping], top_k_sections=2, external=True)
    assert [item["content"] for item in section_result["results"]] == ["Alpha", "Beta"]
    assert all(item["wiki_slug"] == "alice" for item in section_result["results"])
    assert all(item["kind"] == "wiki_section" for item in section_result["results"])
    assert section_result["results"][0]["wiki_page_id"]
    assert "page_content" in section_result["results"][0]

    narrow_service = WikiVectorSearchService(
        LocalWikiIndexConfig(min_similarity=0, top_k_sections=1, top_k_pages=1),
        FakeEmbedding(),
        WikiVectorCache(store, revision_check_seconds=0),
    )
    assert [item["wiki_slug"] for item in (await narrow_service.search("alice", [mapping]))["results"]] == ["alice"]


@pytest.mark.asyncio
async def test_corrupt_blob_marks_index_for_rebuild() -> None:
    store = MemoryLlmWikiIndexStore()
    await store.try_acquire_sync_lease("kb", owner="w", sync_id="s", lease_seconds=60, fingerprint="fp")
    page = normalize_wiki_page("kb", {"slug": "broken", "title": "Broken", "content": "Broken", "status": "published"})
    blob, _ = encode_vector([1.0, 0.0], dimensions=2)
    saved = await store.replace_page_vectors(
        page=page,
        vectors=[{"section_index": 0, "heading": None, "section_content": "Broken", "content_hash": "x", "embedding_dimensions": 2, "vector_blob": blob}],
        fingerprint="fp",
        sync_id="s",
    )
    store.vectors[saved["id"]][0]["vector_blob"] = b"broken"
    service = WikiVectorSearchService(
        LocalWikiIndexConfig(min_similarity=0, top_k_sections=10, top_k_pages=8),
        FakeEmbedding(),
        WikiVectorCache(store, revision_check_seconds=0),
    )

    with pytest.raises(WikiIndexError) as exc_info:
        await service.search("broken", [{"id": "kb", "name": "Broken"}])
    assert exc_info.value.code == "WIKI_INDEX_REBUILDING"
    assert store.states["kb"]["state"] == "rebuild_required"