from techdocs_graphrag.answering import NO_CONTEXT_ANSWER, GraphRAGAssistant
from techdocs_graphrag.extraction import GraphExtractor
from techdocs_graphrag.indexing import IndexingPipeline
from techdocs_graphrag.ingestion import ParagraphSplitter, SourceDocument
from techdocs_graphrag.models import Chunk, GraphFact, RetrievalMode
from techdocs_graphrag.retrieval import HybridRetriever
from techdocs_graphrag.visualization import facts_to_dot

from .conftest import FakeKnowledgeGraph, FakeLanguageModel, retrieved_chunk

EXTRACTION_PAYLOAD = """
{
  "entities": [{"name": "StateGraph", "type": "Component"}],
  "relations": [{"source": "StateGraph", "relation": "uses", "target": "Channel", "evidence": "x"}]
}
"""


def build_retriever(llm: FakeLanguageModel, store: FakeKnowledgeGraph) -> HybridRetriever:
    return HybridRetriever(llm=llm, store=store, top_k=5, max_hops=2, max_facts=10)


def test_graph_mode_seeds_expansion_with_lexical_then_mentioned_entities(
    fake_llm: FakeLanguageModel,
) -> None:
    store = FakeKnowledgeGraph(
        chunks=[
            retrieved_chunk("a.md", ["checkpointer", "stategraph"]),
            retrieved_chunk("b.md", ["checkpointer"]),
        ],
        lexical_matches=["stategraph"],
    )

    context = build_retriever(fake_llm, store).retrieve("question", RetrievalMode.GRAPH)

    assert store.expanded_with == ["stategraph", "checkpointer"]
    assert context.seed_entities == ["stategraph", "checkpointer"]


def test_vector_mode_never_touches_the_graph(fake_llm: FakeLanguageModel) -> None:
    store = FakeKnowledgeGraph(chunks=[retrieved_chunk("a.md", ["stategraph"])])

    context = build_retriever(fake_llm, store).retrieve("question", RetrievalMode.VECTOR)

    assert store.expanded_with is None
    assert context.facts == []


def test_assistant_does_not_call_the_llm_without_context(fake_llm: FakeLanguageModel) -> None:
    assistant = GraphRAGAssistant(retriever=build_retriever(fake_llm, FakeKnowledgeGraph()), llm=fake_llm)

    answer = assistant.ask("question")

    assert answer.text == NO_CONTEXT_ANSWER
    assert fake_llm.completions == []


def test_assistant_sends_graph_facts_to_the_llm(fake_llm: FakeLanguageModel) -> None:
    store = FakeKnowledgeGraph(
        chunks=[retrieved_chunk("a.md", ["stategraph"])],
        facts=[GraphFact(source="StateGraph", relation="USES", target="Channel", documents=["a.md"])],
    )
    assistant = GraphRAGAssistant(retriever=build_retriever(fake_llm, store), llm=fake_llm)

    answer = assistant.ask("Comment fonctionne StateGraph ?")

    _, user_prompt = fake_llm.completions[0]
    assert "StateGraph --USES--> Channel" in user_prompt
    assert answer.context.sources == ["a.md"]


def test_extractor_returns_empty_graph_on_invalid_json() -> None:
    chunk = Chunk(id="1", document="a.md", position=0, text="texte")
    assert GraphExtractor(FakeLanguageModel("pas du json")).extract(chunk).entities == []


def test_indexing_pipeline_reports_extracted_graph() -> None:
    llm = FakeLanguageModel(EXTRACTION_PAYLOAD)
    store = FakeKnowledgeGraph()
    pipeline = IndexingPipeline(
        splitter=ParagraphSplitter(chunk_size=50, chunk_overlap=0),
        llm=llm,
        extractor=GraphExtractor(llm),
        store=store,
        workers=2,
    )
    progress: list[int] = []

    report = pipeline.run(
        [SourceDocument("a.md", "Premier paragraphe.\n\nSecond paragraphe assez long pour couper.")],
        on_progress=lambda done, _: progress.append(done),
    )

    assert (report.chunks, report.entities, report.relations) == (2, 4, 2)
    assert progress == [1, 2]
    assert store.indexes_awaited


def test_dot_output_escapes_quotes() -> None:
    dot = facts_to_dot([GraphFact(source='Say "hi"', relation="USES", target="B")])
    assert '"Say \\"hi\\""' in dot