import os
import numpy as np
import fitz  # PyMuPDF pour extraction PDF
import faiss
import openai
from openai import OpenAI
from sklearn.manifold import TSNE
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from dotenv import load_dotenv

# Charger les variables d'environnement
load_dotenv()
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')

# 📌 Initialisation du client OpenAI
client = OpenAI(api_key='') #todo : mettre la clé en secret et get avec os env
model_embedding = "text-embedding-ada-002"
model_chat = "gpt-4-turbo"

# 📌 Paramètres de segmentation
chunk_size = 256
chunk_overlap = 10

# 📌 Extraction et segmentation des PDF
def extract_and_chunk_pdfs(pdf_folder):
    """Extrait et segmente les textes des PDF en chunks optimisés pour OpenAI."""
    documents = SimpleDirectoryReader(pdf_folder).load_data()
    chunked_docs = [doc.text for doc in documents]
    return chunked_docs

# 📌 Génération des embeddings par batch
def get_embeddings_in_batches(text_chunks, batch_size=5):
    """Génère les embeddings en batch pour éviter les dépassements de tokens."""
    embeddings = []
    for i in range(0, len(text_chunks), batch_size):
        batch = text_chunks[i:i + batch_size]
        response = client.embeddings.create(
            input=batch,
            model=model_embedding
        )
        batch_embeddings = [data.embedding for data in response.data]
        embeddings.extend(batch_embeddings)
    
    return np.array(embeddings).astype('float32')

# 📌 Chargement et embedding des documents
pdf_folder = 'C:/Users/MIPO10053340/OneDrive - Groupe Avril/Bureau/Salon_Agriculture_2024/Micka_API_Call/Docs_pdf/'
chunked_docs = extract_and_chunk_pdfs(pdf_folder)
embeddings = get_embeddings_in_batches(chunked_docs)

# 📌 Indexation des embeddings avec FAISS
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)

# 📌 Récupération des chunks les plus pertinents
def retrieve_relevant_chunks(question, k=5):
    """Recherche les chunks les plus pertinents en fonction de la similarité des embeddings."""
    response = client.embeddings.create(
        input=[question],
        model=model_embedding
    )
    question_embedding = np.array(response.data[0].embedding).astype('float32').reshape(1, -1)
    distances, indices = index.search(question_embedding, k)
    return [chunked_docs[i] for i in indices[0]]

# 📌 Génération de réponse avec OpenAI
def generate_response(context, question):
    """Génère une réponse basée sur le contexte extrait du corpus."""
    messages = [
        {"role": "system", "content": f"Voici des informations contextuelles : {context}"},
        {"role": "user", "content": question}
    ]
    
    response = client.chat.completions.create(
        model=model_chat,
        messages=messages
    )
    return response.choices[0].message.content

# 📌 Exécuter une requête utilisateur
user_question = "Quelles souches de poulet et poules se trouvent dans ce corpus de texte ?"
relevant_chunks = retrieve_relevant_chunks(user_question)
context = "\n".join(relevant_chunks)
answer = generate_response(context, user_question)

# 📊 Affichage de la réponse
print("\n🔹 Réponse OpenAI :")
print(answer)

# 💾 Sauvegarde des résultats
with open("openai_response.txt", "w", encoding="utf-8") as f:
    f.write(f"Question : {user_question}\n")
    f.write(f"Réponse :\n{answer}\n")

print("\n✅ Réponse enregistrée dans 'openai_response.txt'")