import gradio as gr
from rag_pipeline import RAGPipeline
from utils import process_documents
import time

rag = RAGPipeline()

def log_message(msg, logs):
    logs = logs + msg + "\n"
    return logs

def upload_and_index(files, logs):
    logs = log_message("[RAG] بدء معالجة الملفات...", logs)
    all_chunks = []
    for file in files:
        logs = log_message(f"[RAG] معالجة الملف: {file.name}", logs)
        chunks = process_documents(file.name)
        all_chunks.extend(chunks)
        logs = log_message(f"[RAG] تم استخراج {len(chunks)} مقطع من {file.name}", logs)

    logs = log_message(f"[RAG] بناء الفهرس لـ {len(all_chunks)} مقطع...", logs)
    start = time.time()
    rag.build_index(all_chunks)
    duration = time.time() - start
    logs = log_message(f"[RAG] تم بناء الفهرس في {duration:.2f} ثانية.", logs)
    return logs, gr.update(visible=True), gr.update(visible=True)

def answer_question(question, logs):
    logs = log_message(f"[RAG] استلام السؤال: {question}", logs)
    start = time.time()
    answer, sources = rag.answer(question)
    duration = time.time() - start
    logs = log_message(f"[RAG] تم الإجابة في {duration:.2f} ثانية.", logs)
    logs = log_message(f"[RAG] المصادر: {sources}", logs)
    return answer, logs

with gr.Blocks() as demo:
    logs = gr.State("")
    gr.Markdown("# نظام استرجاع المعرفة (RAG)")

    with gr.Row():
        files_input = gr.File(file_types=['.pdf', '.docx', '.txt'], file_count="multiple", label="رفع الملفات")
        upload_btn = gr.Button("رفع وبناء الفهرس")

    logs_output = gr.Textbox(label="سجل العمليات", lines=12, interactive=False, value="")
    question_input = gr.Textbox(label="اكتب سؤالك هنا", visible=False)
    ask_btn = gr.Button("إرسال السؤال", visible=False)
    answer_output = gr.Textbox(label="الإجابة", lines=5)

    upload_btn.click(upload_and_index, inputs=[files_input, logs], outputs=[logs_output, question_input, ask_btn])
    ask_btn.click(answer_question, inputs=[question_input, logs], outputs=[answer_output, logs_output])

demo.launch()