import gradio as gr
from huggingface_hub import InferenceClient
from transformers import AutoTokenizer, AutoModel
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
import fitz  # PyMuPDF


def get_hf_models():
    return ["Qwen/Qwen2.5-3B-Instruct", "HuggingFaceH4/zephyr-7b-beta", "mistralai/Mistral-7B-Instruct-v0.1"]

def extract_text_from_pdf(pdf_path):
    text = ""
    with fitz.open(pdf_path) as doc:
        for page in doc:
            text += page.get_text()
    return text

def manual_rag(query, context, client):
    prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
    response = client.text_generation(prompt, max_new_tokens=512)
    return response

def classic_rag(query, pdf_path, client, embedder):
    text = extract_text_from_pdf(pdf_path)
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = text_splitter.split_text(text)
    embeddings = HuggingFaceEmbeddings(model_name=embedder)
    db = FAISS.from_texts(chunks, embeddings)
    docs = db.similarity_search(query, k=3)
    context = " ".join([doc.page_content for doc in docs])
    response = manual_rag(query, context, client)
    return response, context

def no_rag(query, client):
    response = client.text_generation(query, max_new_tokens=512)
    return response

def process_query(query, pdf_path, llm_choice, embedder_choice, use_manual_rag):
    client = InferenceClient(llm_choice)
    no_rag_response = no_rag(query, client)

    if pdf_path is None:
        return no_rag_response, "RAG non utilisé (pas de fichier PDF)", "RAG non utilisé (pas de fichier PDF)", "Pas de fichier PDF fourni", "Pas de contexte extrait"

    full_text = extract_text_from_pdf(pdf_path)

    # RAG manuel seulement si choisi
    if use_manual_rag == "Oui":
        manual_rag_response = manual_rag(query, full_text, client)
    else:
        manual_rag_response = "RAG manuel non utilisé"

    classic_rag_response, classic_rag_context = classic_rag(query, pdf_path, client, embedder_choice)

    return no_rag_response, manual_rag_response, classic_rag_response, full_text, classic_rag_context


iface = gr.Interface(
    fn=process_query,
    inputs=[
        gr.Textbox(label="Votre question"),
        gr.File(label="Chargez un nouveau PDF"),
        gr.Dropdown(choices=get_hf_models(), label="Choisissez le LLM", value="HuggingFaceH4/zephyr-7b-beta"),
        gr.Dropdown(choices=["sentence-transformers/all-MiniLM-L6-v2", "nomic-ai/nomic-embed-text-v1.5"],
                    label="Choisissez l'Embedder", value="sentence-transformers/all-MiniLM-L6-v2"),
        gr.Dropdown(choices=["Oui", "Non"], label="Utiliser RAG manuel ?", value="Non")  # Ajout de la combobox pour choisir RAG manuel
    ],
    outputs=[
        gr.Textbox(label="Réponse sans RAG"),
        gr.Textbox(label="Réponse avec RAG manuel"),
        gr.Textbox(label="Réponse avec RAG classique"),
        gr.Textbox(label="Texte complet du PDF (pour RAG manuel)", lines=10),
        gr.Textbox(label="Contexte extrait (pour RAG classique)", lines=10)
    ],
    title="Tutoriel RAG - Comparaison des méthodes",
    description="Posez une question sur le contenu d'un PDF et comparez les réponses obtenues avec différentes méthodes de RAG.",
    theme="default"
)

if __name__ == "__main__":
    iface.launch()