import os
import time
from dotenv import load_dotenv
from pinecone import Pinecone, ServerlessSpec
from langsmith import Client
from langchain_mistralai.chat_models import ChatMistralAI
from neo4j import GraphDatabase
from sentence_transformers import SentenceTransformer
from pinecone_text.sparse import BM25Encoder
from langchain.embeddings import HuggingFaceEmbeddings

# Charger les variables d'environnement
load_dotenv()

# Configuratio Pinecone
pinecone_api_key = os.getenv("PINECONE_API_KEY")
if not pinecone_api_key:
    raise ValueError("PINECONE_API_KEY n'est pas défini dans .env")

# Création du client Pinecone
pc = Pinecone(api_key=pinecone_api_key)

# Configuration des index
sparse_index_name = "confession-muette-sparse"
dense_index_name = "confession-muette-dense"
judge_index_name = "confession-muette-judge"

# LangSmith configuration
langsmith_endpoint = os.getenv("LANGSMITH_ENDPOINT")
langsmith_project = os.getenv("LANGSMITH_PROJECT")
langsmith_api_key = os.getenv("LANGSMITH_API_KEY")
client = Client(
    api_url=langsmith_endpoint,
    api_key=langsmith_api_key,
)
dataset_id = os.getenv("DATASET_UUID")

# Initialize Neo4j connection
neo4j_uri = os.getenv("NEO4J_URI")
neo4j_username = os.getenv("NEO4J_USERNAME")
neo4j_password = os.getenv("NEO4J_PASSWORD")
neo4j_driver = GraphDatabase.driver(
    neo4j_uri,
    auth=(neo4j_username, neo4j_password)
)

# Mistral AI configuration
mistral_api_key = os.getenv("MISTRAL_API_KEY")
llm = ChatMistralAI(
    model="mistral-large-latest",
    api_key=mistral_api_key,
    temperature=0,
    stream=False,
    verbose=True
)

# Initialiser les modèles et encodeurs
model = SentenceTransformer("intfloat/multilingual-e5-large")
sparse_encoder = BM25Encoder().default()
embeddings = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")

# Vérifier si les index existent
existing_indexes = pc.list_indexes()
all_names = [idx["name"] for idx in existing_indexes]

# Créer l'index sparse s'il n'existe pas
if sparse_index_name not in all_names:
    print(f"L'index '{sparse_index_name}' n'existe pas. Création en cours...")
    try:
        pc.create_index(
            name=sparse_index_name,
            metric="dotproduct",
            dimension=1024,
            spec=ServerlessSpec(cloud="aws", region="us-east-1"),  # Modifier la région si besoin
        )
        print(f"Index '{sparse_index_name}' créé avec succès.")
    except Exception as e:
        print("Erreur lors de la création de l'index sparse:", e)

# Créer l'index dense s'il n'existe pas
if dense_index_name not in all_names:
    print(f"L'index '{dense_index_name}' n'existe pas. Création en cours...")
    try:
        pc.create_index(
            name=dense_index_name,
            dimension=1024,
            metric="cosine",
            spec=ServerlessSpec(cloud="aws", region="us-east-1"),
        )
        
        print(f"Index '{dense_index_name}' créé avec succès.")
    except Exception as e:
        print("Erreur lors de la création de l'index dense:", e)

# Créer l'index judge s'il n'existe pas
if judge_index_name not in all_names:
    print(f"L'index '{judge_index_name}' n'existe pas. Création en cours...")
    try:
        pc.create_index(
            name=judge_index_name,
            metric="dotproduct",
            dimension=1024,
            spec=ServerlessSpec(cloud="aws", region="us-east-1"),  
        )
        print(f"Index '{judge_index_name}' créé avec succès.")
    except Exception as e:
        print("Erreur lors de la création de l'index sparse:", e)

print(f"Les index '{sparse_index_name}' , '{dense_index_name}' et '{judge_index_name}' sont prêts.")

# Connexion aux index
sparse_index = pc.Index(sparse_index_name)
dense_index = pc.Index(dense_index_name)
judge_index = pc.Index(judge_index_name)
print(f"Connexion aux index '{sparse_index_name}' , '{dense_index_name}' et '{judge_index_name}' réussie.")