from collections.abc import Sequence
from pathlib import Path

import streamlit as st
from neo4j.exceptions import DriverError, Neo4jError
from pydantic import ValidationError

from techdocs_graphrag import (
    Answer,
    RetrievalMode,
    Services,
    SourceDocument,
    build_services,
    load_documents,
)
from techdocs_graphrag.visualization import facts_to_dot

SAMPLE_CORPUS = Path(__file__).parent / "data"

MODE_LABELS = {
    RetrievalMode.GRAPH: "GraphRAG",
    RetrievalMode.VECTOR: "RAG vectoriel",
}

DISPLAY_MODES: dict[str, tuple[RetrievalMode, ...]] = {
    "GraphRAG": (RetrievalMode.GRAPH,),
    "RAG vectoriel": (RetrievalMode.VECTOR,),
    "Comparaison côte à côte": (RetrievalMode.VECTOR, RetrievalMode.GRAPH),
}

EXAMPLE_QUESTIONS = (
    "Comment persister l'état d'un workflow StateGraph ?",
    "Quel est le rôle des reducers dans les channels ?",
    "Comment reprendre un workflow après une interruption humaine ?",
)

st.set_page_config(page_title="TechDocs GraphRAG", page_icon="🕸️", layout="wide")


@st.cache_resource(show_spinner="Connexion à Neo4j et OpenAI…")
def get_services() -> Services:
    services = build_services()
    services.store.verify_connectivity()
    services.store.ensure_schema()
    return services


def load_services() -> Services:
    try:
        return get_services()
    except ValidationError as error:
        missing = ", ".join(sorted({str(item["loc"][0]).upper() for item in error.errors()}))
        st.error(f"Configuration incomplète. Variables à définir dans les secrets : {missing}")
    except (Neo4jError, DriverError) as error:
        st.error(f"Connexion à Neo4j impossible : {error}")
    st.stop()


def index_documents(services: Services, documents: Sequence[SourceDocument]) -> None:
    progress = st.progress(0.0, text="Calcul des embeddings…")

    def report_progress(done: int, total: int) -> None:
        progress.progress(done / total, text=f"Extraction du graphe : {done}/{total} chunks")

    report = services.indexer.run(documents, on_progress=report_progress)
    progress.empty()
    st.success(
        f"{report.documents} documents · {report.chunks} chunks · "
        f"{report.entities} entités · {report.relations} relations"
    )


def render_sidebar(services: Services) -> tuple[RetrievalMode, ...]:
    with st.sidebar:
        st.header("Base de connaissances")

        uploads = st.file_uploader(
            "Ajouter des documents",
            type=["md", "markdown", "txt"],
            accept_multiple_files=True,
        )
        if st.button("Indexer les documents ajoutés", disabled=not uploads):
            index_documents(
                services,
                [SourceDocument(name=file.name, text=file.getvalue().decode("utf-8")) for file in uploads],
            )
        if st.button("Indexer le corpus d'exemple"):
            index_documents(services, load_documents(SAMPLE_CORPUS))

        st.divider()
        display_mode = st.radio("Mode de recherche", options=list(DISPLAY_MODES))

        st.divider()
        render_stats(services)
        if st.button("Effacer la conversation"):
            st.session_state.turns = []
            st.rerun()

    return DISPLAY_MODES[display_mode]


def render_stats(services: Services) -> None:
    stats = services.store.stats()
    first, second = st.columns(2)
    first.metric("Documents", stats["documents"])
    second.metric("Chunks", stats["chunks"])
    first.metric("Entités", stats["entities"])
    second.metric("Relations", stats["relations"])


def render_answer(answer: Answer) -> None:
    st.markdown(answer.text)
    context = answer.context
    if context.sources:
        st.caption("Sources : " + " · ".join(context.sources))

    with st.expander(f"Contexte · {len(context.chunks)} extraits · {len(context.facts)} relations"):
        graph_tab, facts_tab, chunks_tab = st.tabs(["Sous-graphe", "Relations", "Extraits"])
        with graph_tab:
            if context.facts:
                st.graphviz_chart(facts_to_dot(context.facts, context.seed_entities))
            else:
                st.info("Aucune relation du graphe n'a été utilisée.")
        with facts_tab:
            if context.facts:
                st.dataframe([fact.model_dump() for fact in context.facts], hide_index=True)
            else:
                st.info("Aucune relation.")
        with chunks_tab:
            for chunk in context.chunks:
                st.markdown(f"**{chunk.document}** · chunk {chunk.position} · score {chunk.score:.3f}")
                st.text(chunk.text)


def render_turn(answers: Sequence[Answer]) -> None:
    with st.chat_message("user"):
        st.markdown(answers[0].question)
    with st.chat_message("assistant"):
        for column, answer in zip(st.columns(len(answers)), answers, strict=True):
            with column:
                if len(answers) > 1:
                    st.subheader(MODE_LABELS[answer.mode])
                render_answer(answer)


def pick_example_question() -> str | None:
    st.markdown("##### Exemples de questions")
    columns = st.columns(len(EXAMPLE_QUESTIONS))
    clicked = [
        question
        for column, question in zip(columns, EXAMPLE_QUESTIONS, strict=True)
        if column.button(question)
    ]
    return clicked[0] if clicked else None


def main() -> None:
    services = load_services()
    modes = render_sidebar(services)

    st.title("🕸️ TechDocs GraphRAG")
    st.caption(
        "Assistant de documentation technique : recherche vectorielle et graphe de connaissances "
        "Neo4j pour des réponses multi-hop sourcées."
    )

    turns: list[list[Answer]] = st.session_state.setdefault("turns", [])
    for answers in turns:
        render_turn(answers)

    typed_question = st.chat_input("Pose une question sur la documentation…")
    question = typed_question or (None if turns else pick_example_question())
    if not question:
        return

    with st.spinner("Recherche dans le graphe et les documents…"):
        answers = [services.assistant.ask(question, mode) for mode in modes]
    turns.append(answers)
    render_turn(answers)


main()