from collections.abc import Sequence

import pytest

from techdocs_graphrag.models import Chunk, GraphExtraction, GraphFact, RetrievedChunk


class FakeLanguageModel:
    def __init__(self, completion: str = "réponse") -> None:
        self.completion = completion
        self.completions: list[tuple[str, str]] = []

    def complete(self, system_prompt: str, user_prompt: str, *, json_mode: bool = False) -> str:
        self.completions.append((system_prompt, user_prompt))
        return self.completion

    def embed(self, texts: Sequence[str]) -> list[list[float]]:
        return [[float(len(text))] for text in texts]


class FakeKnowledgeGraph:
    def __init__(
        self,
        chunks: list[RetrievedChunk] | None = None,
        lexical_matches: list[str] | None = None,
        facts: list[GraphFact] | None = None,
    ) -> None:
        self.chunks = chunks or []
        self.lexical_matches = lexical_matches or []
        self.facts = facts or []
        self.expanded_with: list[str] | None = None
        self.extractions: list[tuple[Chunk, GraphExtraction]] = []
        self.indexes_awaited = False

    def ensure_schema(self) -> None:
        pass

    def await_indexes(self) -> None:
        self.indexes_awaited = True

    def upsert_chunks(self, chunks: Sequence[Chunk], embeddings: Sequence[list[float]]) -> None:
        pass

    def upsert_extraction(self, chunk: Chunk, extraction: GraphExtraction) -> None:
        self.extractions.append((chunk, extraction))

    def vector_search(self, embedding: list[float], top_k: int) -> list[RetrievedChunk]:
        return self.chunks[:top_k]

    def search_entities(self, text: str, limit: int) -> list[str]:
        return self.lexical_matches[:limit]

    def expand(self, seed_keys: Sequence[str], max_hops: int, limit: int) -> list[GraphFact]:
        self.expanded_with = list(seed_keys)
        return self.facts[:limit]


def retrieved_chunk(document: str, entity_keys: list[str]) -> RetrievedChunk:
    return RetrievedChunk(
        id=f"{document}-0",
        document=document,
        position=0,
        text=f"contenu de {document}",
        score=0.9,
        entity_keys=entity_keys,
    )


@pytest.fixture
def fake_llm() -> FakeLanguageModel:
    return FakeLanguageModel()