from typing import Protocol

from zbt.domains.knowledge.models import KnowledgeChunk, KnowledgeSearchHit


class KnowledgeIndex(Protocol):
    def replace_document(
        self,
        document_id: str,
        chunks: list[KnowledgeChunk],
    ) -> None: ...

    def search(
        self,
        query: str,
        *,
        document_ids: set[str],
        limit: int,
    ) -> list[KnowledgeSearchHit]: ...


class InMemoryKnowledgeIndex:
    def __init__(self) -> None:
        self._chunks_by_document: dict[str, list[KnowledgeChunk]] = {}

    def replace_document(
        self,
        document_id: str,
        chunks: list[KnowledgeChunk],
    ) -> None:
        self._chunks_by_document[document_id] = list(chunks)

    def search(
        self,
        query: str,
        *,
        document_ids: set[str],
        limit: int,
    ) -> list[KnowledgeSearchHit]:
        query_terms = {character for character in query if character.strip()}
        hits: list[KnowledgeSearchHit] = []
        for document_id in document_ids:
            for chunk in self._chunks_by_document.get(document_id, []):
                content_terms = set(chunk.content)
                overlap = len(query_terms & content_terms)
                if overlap == 0:
                    continue
                hits.append(
                    KnowledgeSearchHit(
                        chunk_id=chunk.id,
                        document_id=chunk.document_id,
                        document_version=chunk.document_version,
                        ordinal=chunk.ordinal,
                        title=chunk.title,
                        content=chunk.content,
                        document_type=chunk.document_type,
                        source_name=chunk.source_name,
                        product_code=chunk.product_code,
                        score=overlap / max(len(query_terms), 1),
                    )
                )
        return sorted(hits, key=lambda item: item.score, reverse=True)[:limit]
