from typing import TypedDict, Annotated, Sequence
from langchain_core.messages import BaseMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from config import llm, client, langsmith_project
from pinecone_utilsB import *  
from typing import TypedDict, Sequence, List, Dict, Optional, Annotated


class GraphState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]
    query: str
    relevant_docs: List[Dict[str, Optional[Dict[str, float]]]]
    response: str
    k: int 
    alpha: float
    similarity_threshold: float


def retrieve_combined(state: GraphState) -> dict:
    """Récupération hybride : Pinecone (sémantique) + BM25 (mots-clés)."""

    relevant_docs = hybrid_search(
        state["query"], 
        alpha=state.get("alpha"), 
        k=state.get("k"), 
        similarity_threshold=state.get("similarity_threshold")
    )
    
    return {"relevant_docs": relevant_docs}

def generate_response(state: GraphState) -> dict:
    """Génération de réponse en combinant informations sémantiques et mots-clés."""
    context = "\n\n".join(doc["text"] for doc in state["relevant_docs"])
    prompt = f"""
    Vous êtes un expert en analyse de texte.
    Votre tâche consiste à répondre à la question de l'utilisateur en utilisant les informations pertinentes fournies. 
    Intégrez à la fois les éléments sémantiques (contexte, sens) et les mots-clés exacts pour fournir une réponse fluide et cohérente.
    
    **Instructions supplémentaires** :
    - Utilisez les mots-clés pertinents de manière naturelle dans votre réponse.
    - Expliquez les concepts en vous appuyant sur le contexte sémantique.
    - Ne mentionnez pas explicitement les termes "recherche sémantique" ou "recherche par mots-clés".
    - **Si la réponse doit être une liste retournez chaque élément sur **une nouvelle ligne**.
    - Format correct attendu :
      1. Élément 1
      2. Élément 2
      3. Élément 3
    
    **Informations pertinentes trouvées** :
    {context}
    
    **Question de l'utilisateur** :
    {state["query"]}
    
    **Réponse :**
    [Fournissez une réponse cohérente qui intègre à la fois les éléments sémantiques et les mots-clés pertinents sans les distinguer explicitement.]
    """
    response = llm.invoke(prompt)
    return {"response": response.content}


def post_process_response(state: GraphState) -> dict:
    """Nettoie et valide la réponse."""
    response = state["response"].strip()
    
    # Vérifier si la réponse est pertinente
    if not response or response.lower() in ["je ne sais pas", "i don't know"]:
        response = "Désolé, je n'ai pas trouvé d'informations pertinentes pour votre question."
    
    return {"response": response}

# Construction du graphe
graph_builder = StateGraph(GraphState)
graph_builder.add_node("retrieve", retrieve_combined)
graph_builder.add_node("generate", generate_response)
graph_builder.add_node("post_process", post_process_response)

graph_builder.set_entry_point("retrieve")
graph_builder.add_edge("retrieve", "generate")
graph_builder.add_edge("generate", "post_process")
graph_builder.add_edge("post_process", END)

agent = graph_builder.compile()