from collections import Counter
from collections.abc import Sequence

from techdocs_graphrag.models import RetrievalContext, RetrievalMode, RetrievedChunk
from techdocs_graphrag.ports import KnowledgeGraph, LanguageModel

MAX_LEXICAL_SEEDS = 5
MAX_MENTION_SEEDS = 8


class HybridRetriever:
    def __init__(
        self,
        *,
        llm: LanguageModel,
        store: KnowledgeGraph,
        top_k: int,
        max_hops: int,
        max_facts: int,
    ) -> None:
        self._llm = llm
        self._store = store
        self._top_k = top_k
        self._max_hops = max_hops
        self._max_facts = max_facts

    def retrieve(self, question: str, mode: RetrievalMode = RetrievalMode.GRAPH) -> RetrievalContext:
        embedding = self._llm.embed([question])[0]
        chunks = self._store.vector_search(embedding, self._top_k)
        if mode is RetrievalMode.VECTOR:
            return RetrievalContext(chunks=chunks)

        seeds = self._seed_entities(question, chunks)
        facts = self._store.expand(seeds, self._max_hops, self._max_facts)
        return RetrievalContext(chunks=chunks, facts=facts, seed_entities=seeds)

    def _seed_entities(self, question: str, chunks: Sequence[RetrievedChunk]) -> list[str]:
        lexical_matches = self._store.search_entities(question, MAX_LEXICAL_SEEDS)
        mentions = Counter(key for chunk in chunks for key in chunk.entity_keys)
        most_mentioned = [key for key, _ in mentions.most_common(MAX_MENTION_SEEDS)]
        return list(dict.fromkeys([*lexical_matches, *most_mentioned]))