import gradio as gr
from src.utils.search_docs_utils import search_relevant_documents
from src.rag.documents_rag_pipeline import RAGPipeline

# Charger les 25 questions depuis le fichier texte
with open("docs/questions_rag_test.txt", encoding="utf-8") as f:
    questions_list = f.read()

def rag_answer(query, top_k_docs=5, top_k_chunks=5):
    result = search_relevant_documents(query, top_k=top_k_docs)
    docs = result.get('documents', []) if 'documents' in result else result.get('results', [])
    rag_pipeline = RAGPipeline()
    print(f'docs : {docs}')
    top_chunks = rag_pipeline.get_top_docs_chunks_for_query(query, docs)
    answer = rag_pipeline.generate(query, top_chunks)
    # Format docs and chunks for display
    docs_str = '\n\n'.join([f"- {doc.get('Nom du document', doc.get('pdf_title', ''))} ({doc.get('Lien', doc.get('pdf_link', ''))})" for doc in docs])
    chunks_str = '\n---\n'.join([chunk['text'] for chunk in top_chunks])
    return answer, docs_str, chunks_str

with gr.Blocks() as demo:
    gr.Markdown("# RAG QA Demo\nPosez votre question sur les documents marocains !")
    with gr.Row():
        with gr.Column(scale=3):
            question = gr.Textbox(label="Votre question", lines=2)
            with gr.Row():
                top_k_docs = gr.Slider(1, 10, value=5, step=1, label="Top K Documents")
                top_k_chunks = gr.Slider(1, 20, value=10, step=1, label="Top K Chunks")
            submit = gr.Button("Générer la réponse")
            answer = gr.Textbox(label="Réponse générée", lines=5)
            docs = gr.Textbox(label="Documents pertinents", lines=5)
            chunks = gr.Textbox(label="Chunks utilisés", lines=5)
            submit.click(rag_answer, inputs=[question, top_k_docs, top_k_chunks], outputs=[answer, docs, chunks])
        with gr.Column(scale=2):
            gr.Markdown("""### Questions de test (copier-coller)

```text
""" + questions_list + """
```
""")

demo.launch(share=True)