from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from sentence_transformers import SentenceTransformer
import fitz  # PyMuPDF pour extraction du texte PDF
import faiss
import os
import numpy as np
import requests
import gradio as gr

# Configuration de l'API Gemini
GEMINI_API_KEY = "AIzaSyArbgg_p_HlmpgrcjVYemdSJeMCP9OTj3E"
GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent"

# Configuration des embeddings et FAISS
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
model = SentenceTransformer(EMBEDDING_MODEL)
INDEX_PATH = "medical_faiss_index"
documents = []  # Liste pour stocker les textes indexés

# Chargement ou création de l'index FAISS
if os.path.exists(INDEX_PATH):
    index = faiss.read_index(INDEX_PATH)
    print("Index FAISS chargé avec succès.")
else:
    index = faiss.IndexFlatL2(model.get_sentence_embedding_dimension())
    print("Nouvel index FAISS créé.")

# Fonction pour extraire le texte des fichiers PDF
def extract_text_from_pdf(file_content: bytes):
    pdf = fitz.open(stream=file_content, filetype="pdf")
    paragraphs = []
    for page in pdf:
        text = page.get_text()
        if text.strip():
            paragraphs.extend([p.strip() for p in text.split("\n\n") if p.strip()])
    return paragraphs

# Ajouter des documents médicaux à l'index FAISS
def add_medical_reference(file_content: bytes):
    global index, documents
    paragraphs = extract_text_from_pdf(file_content)
    embeddings = model.encode(paragraphs)
    index.add(np.array(embeddings, dtype="float32"))
    documents.extend(paragraphs)
    faiss.write_index(index, INDEX_PATH)

# Recherche dans FAISS pour trouver les documents pertinents
def search_faiss(query, k=5):
    query_embedding = model.encode([query])
    distances, indices = index.search(np.array(query_embedding, dtype="float32"), k)
    results = [documents[i] for i in indices[0] if i < len(documents)]
    return results

# Appel à l'API Gemini
def call_gemini_api(prompt):
    headers = {"Content-Type": "application/json"}
    payload = {
        "contents": [
            {
                "parts": [
                    {"text": prompt}
                ]
            }
        ]
    }
    try:
        response = requests.post(f"{GEMINI_API_URL}?key={GEMINI_API_KEY}", json=payload, headers=headers)
        response.raise_for_status()
        response_json = response.json()
        candidates = response_json.get("candidates", [])
        if candidates:
            return candidates[0]["content"]["parts"][0]["text"]
        return "Pas de réponse disponible depuis Gemini."
    except requests.exceptions.RequestException as e:
        return f"Erreur API Gemini : {str(e)}"

# Ajouter un PDF de référence médicale
def upload_reference(file):
    try:
        # Ouvrir le fichier en mode binaire
        with open(file.name, "rb") as f:
            file_content = f.read()
        add_medical_reference(file_content)
        return "Référence médicale ajoutée avec succès."
    except Exception as e:
        return f"Erreur : {str(e)}"


# Analyser un PDF d'analyse de sang
def analyze_blood_test(file):
    try:
        # Ouvrir le fichier en mode binaire
        with open(file.name, "rb") as f:
            file_content = f.read()
        
        # Extraire le texte des résultats d'analyse
        test_results = extract_text_from_pdf(file_content)

        if not test_results:
            return "Aucun texte valide extrait du fichier d'analyse."

        # Combine tous les résultats d'analyse dans un seul texte
        query = "\n".join(test_results)

        # Recherche dans l'index FAISS
        relevant_docs = search_faiss(query, k=5)
        context = "\n".join(relevant_docs)

        # Enrichir le prompt avec les informations pertinentes
        enriched_prompt = f"Voici les résultats d'analyse :\n{query}\n\nContexte pertinent :\n{context}"
        gemini_response = call_gemini_api(enriched_prompt)

        return {
            "Réponse générée": gemini_response,
            "Documents pertinents": relevant_docs
        }
    except Exception as e:
        return f"Erreur : {str(e)}"


# Interface Gradio
def gradio_upload_reference(file):
    return upload_reference(file)

def gradio_analyze_blood_test(file):
    response = analyze_blood_test(file)
    if isinstance(response, dict):
        return f"Réponse générée :\n{response['Réponse générée']}\n\nDocuments pertinents :\n{response['Documents pertinents']}"
    return response

# Lancer l'application Gradio
with gr.Blocks() as demo:
    gr.Markdown("## Analyse Médicale avec RAG et Gemini")
    
    with gr.Tab("Ajouter Références Médicales"):
        ref_file = gr.File(label="Téléchargez un fichier PDF de référence médicale")
        ref_output = gr.Textbox(label="Résultat")
        ref_button = gr.Button("Ajouter")
        ref_button.click(gradio_upload_reference, inputs=ref_file, outputs=ref_output)
    
    with gr.Tab("Analyser un Résultat d'Analyse"):
        test_file = gr.File(label="Téléchargez un fichier PDF d'analyse de sang")
        analysis_output = gr.Textbox(label="Résultat")
        analyze_button = gr.Button("Analyser")
        analyze_button.click(gradio_analyze_blood_test, inputs=test_file, outputs=analysis_output)

demo.launch()