import os
import gradio as gr
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA
from langchain.llms import HuggingFaceHub
from langchain.prompts import PromptTemplate
from huggingface_hub import login

# اختياري: تسجيل الدخول إذا كنت تستخدم مفتاح API
# login(token="your_huggingface_token")

def process_pdf_and_answer(pdf_path, question):
    # تحميل ملف PDF
    loader = PyPDFLoader(pdf_path)
    pages = loader.load()

    # تقسيم النصوص
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    texts = text_splitter.split_documents(pages)

    # التضمين Embeddings
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
    vectorstore = Chroma.from_documents(texts, embedding=embeddings)

    # إعداد LLM
    llm = HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta", model_kwargs={"temperature": 0.1, "max_new_tokens": 512})

    # إعداد RAG
    qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=vectorstore.as_retriever(), return_source_documents=True)

    # تنفيذ السؤال
    result = qa_chain({"query": question})
    answer = result["result"]
    return answer

# واجهة Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 مساعد PDF الذكي")
    
    with gr.Row():
        file_input = gr.File(label="📄 ارفع ملف PDF", type="filepath", file_types=[".pdf"])
    
    question_input = gr.Textbox(label="❓ اكتب سؤالك هنا", placeholder="ما هو محتوى الفصل الأول؟")
    output = gr.Textbox(label="📝 الإجابة", lines=10)

    submit_btn = gr.Button("🔍 استخرج الإجابة")

    def handle_submit(file, question):
        if file is None or question.strip() == "":
            return "يرجى رفع ملف PDF وكتابة سؤال."
        return process_pdf_and_answer(file, question)

    submit_btn.click(handle_submit, inputs=[file_input, question_input], outputs=output)

# لتشغيل التطبيق
if __name__ == "__main__":
    demo.launch()