from langchain.schema import HumanMessage 
import json
import streamlit as st
from config import neo4j_driver, sparse_index as indexB, llm
from pinecone_utilsB import hybrid_search
from neo4j.exceptions import CypherSyntaxError
import re
import logging

logging.basicConfig(level=logging.INFO)  
logger = logging.getLogger(__name__)

def extract_cypher_query(llm_output):
    """
    Extrait la requête Cypher valide de la sortie du LLM.
    """
    # Utiliser une expression régulière pour trouver la requête Cypher
    cypher_pattern = r"(MATCH|CREATE|MERGE|RETURN|WITH|DELETE|SET|REMOVE|UNWIND|CALL)[\s\S]*?;"
    match = re.search(cypher_pattern, llm_output)
    
    if match:
        return match.group(0)  # Retourne la requête Cypher trouvée
    else:
        # Si aucune requête n'est trouvée, essayer de retourner la dernière ligne
        lines = llm_output.split("\n")
        for line in reversed(lines):
            if line.strip().endswith(";"):
                return line.strip()
        raise ValueError("Aucune requête Cypher valide trouvée dans la sortie du LLM.")

# Function to generate Cypher queries using an LLM
def generate_cypher_query(user_query):
    """Génère une requête Cypher à l'aide du LLM et extrait la partie valide."""
    # Liste des nœuds et relations autorisés
    allowed_relationships = [
    "CONNAÎT", "SITUE_DANS", "FAIT_PARTIE_DE", "SE_PRODUIT_PENDANT", "IMPLIQUE",
    "S'OPPOSE_À", "CRÉÉ_PAR", "INSPIRÉ_PAR", "REPRÉSENTE", "TRANSFORME",
    "A_PARTICIPÉ_À", "A_CAUSÉ", "POSSEDE", "A_MODIFIÉ", "A_SURVÉCU_À", 
    "A_ÉTÉ_CRÉÉ_PENDANT", "A_ÉTÉ_TRANSMIS_PENDANT", "EST_ISSU_DE", "APPARTIENT_À"
    ]
    allowed_nodes = [
        "Personnage", "Organisation", "Lieu","Pays", "Groupe",
        "Événement", "Objet", "Théme",
    ]

    # Construire le prompt avec des instructions claires
    prompt = f"""
    Vous êtes un générateur de requêtes Cypher pour une base de données Neo4j. 
    Étant donné une demande de l'utilisateur, générez une requête Cypher correspondante.
    **Instructions supplémentaires** :
    - Utilisez `DISTINCT` pour éviter les répétitions.
    - Optimisez la requête pour éviter les chemins redondants.
    - Tous les nœuds ont uniquement la propriété `id`. Utilisez `id` pour filtrer ou retourner des valeurs.
    - Utilisez uniquement les nœuds et relations autorisés suivants :
    **Nœuds autorisés** :
    {", ".join(allowed_nodes)}
    **Relations autorisées** :
    {", ".join(allowed_relationships)}
    **Exemples de requêtes** :
    - Demande : "Trouver tous les lieux mentionnés dans une histoire."
      Requête : MATCH (n:Personnage) RETURN n.id;
    - Demande : "Lister tous les personnages qui connaissent Zéphyrine."
      Requête : MATCH (p1:Personnage)-[:CONNAÎT]->(p2:Personnage {{name: "Zéphyrine"}}) RETURN DISTINCT p1.name
    **Format de la réponse** :
    - Ne fournissez **que la requête Cypher**, sans explications ni commentaires.
    - La requête doit commencer par `MATCH`, `CREATE`, `MERGE`, `RETURN`, etc.
    - La requête doit se terminer par un point-virgule (`;`).
    - Si vous ne pouvez pas générer de requête, retournez une chaîne vide.
    **Demande de l'utilisateur** :
    {user_query}
    **Requête Cypher** :
    """
    
    # Générer la réponse avec le LLM
    response = llm.invoke(prompt)
    logger.debug(f"Sortie brute du LLM : {response.content}")

    # Extraire la requête Cypher valide
    try:
        cypher_query = extract_cypher_query(response.content)
        return cypher_query
    except ValueError as e:
        print(f"Erreur lors de l'extraction de la requête Cypher : {e}")
        return None

# Function to execute Cypher queries in Neo4j
def execute_cypher_query(cypher_query):
    """Exécute une requête Cypher dans Neo4j et retourne les résultats."""
    if not cypher_query:
        return []
    
    try:
        with neo4j_driver.session() as session:
            result = session.run(cypher_query)
            return [record for record in result]
    except CypherSyntaxError as e:
        print(f"Erreur de syntaxe Cypher : {e}")
        return []

# Function to merge hybrid search and Neo4j results
def merge_results(hybrid_results, neo4j_results):
    """Merge results from hybrid search and Neo4j into a unified list."""
    combined_results = []

    # Add hybrid search results
    for doc in hybrid_results:
        combined_results.append({
            "type": "text",
            "content": doc["text"],  # Utilisez la clé "text" du document hybride
            "score": 1.0  
        })

    # Add Neo4j search results
    for record in neo4j_results:
        combined_results.append({
            "type": "entity",
            "content": str(record),  # Convertir en chaîne de caractères si nécessaire
            "score": 0.80  # Score par défaut pour les résultats Neo4j
        })

    # Sort combined results by score (or any other criteria)
    combined_results.sort(key=lambda x: x["score"], reverse=True)

    return combined_results

# Unified search function
def unified_search(query, k, alpha, similarity_threshold):
    """Perform a unified search combining hybrid search and Neo4j knowledge graph search."""
    # Step 1: Perform hybrid search
    hybrid_results = hybrid_search(query, alpha=alpha, k=k, similarity_threshold=similarity_threshold)
    #print(f"résultat recherche hybride: {hybrid_results}")

    # Step 2: Generate and execute Cypher query
    cypher_query = generate_cypher_query(query)
    print(f"Generated Cypher query: {cypher_query}")


    neo4j_results = execute_cypher_query(cypher_query)
    logger.debug(f"Résultats Neo4j AVANT filtration : {neo4j_results}")


    # Step 3: Merge results
    final_results = merge_results(hybrid_results, neo4j_results)

    return final_results