import os
import gradio as gr
from langchain_community.llms import CTransformers
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA

# إعداد النموذج المحلي (تأكد من وجود ملف GGUF داخل مجلد models)
llm = CTransformers(
    model="models/mistral-7b-instruct-v0.2.Q4_K_M.gguf",
    model_type="mistral",
    config={"max_new_tokens": 512, "temperature": 0.5}
)

# إعداد نموذج التضمين
embedding_model = HuggingFaceEmbeddings(
    model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)

# تحميل ومعالجة الملفات
def process_pdf(pdf_file):
    loader = PyPDFLoader(pdf_file.name)
    documents = loader.load()

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50
    )
    docs = text_splitter.split_documents(documents)

    vectordb = Chroma.from_documents(docs, embedding_model, persist_directory="chroma_db")
    vectordb.persist()
    return vectordb

# تهيئة النظام عند تحميل PDF
vectordb = None

def upload_file(file):
    global vectordb
    vectordb = process_pdf(file)
    return "📚 تم تحميل الملف بنجاح ويمكنك الآن طرح الأسئلة."

# الإجابة عن الأسئلة
def answer_question(question):
    global vectordb
    if not vectordb:
        return "❗ يرجى رفع ملف PDF أولًا."

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=vectordb.as_retriever(search_kwargs={"k": 3}),
        return_source_documents=True
    )

    result = qa_chain(question)
    answer = result['result']
    return f"💬 الإجابة:\n\n{answer}"

# واجهة Gradio
with gr.Blocks(title="Smart PDF Assistant", theme=gr.themes.Soft()) as demo:
    gr.Markdown("## 🤖 مساعد PDF الذكي - نظام عربي للإجابة من الملفات بدون API")
    
    with gr.Row():
        pdf_input = gr.File(label="📄 حمّل ملف PDF", file_types=[".pdf"])
        upload_btn = gr.Button("🔁 تحميل ومعالجة الملف")
        upload_output = gr.Textbox(label="الحالة", interactive=False)

    with gr.Row():
        question_input = gr.Textbox(label="❓ اكتب سؤالك هنا", placeholder="ما هو موضوع الفصل الأول؟", lines=2)
        answer_output = gr.Textbox(label="💡 الإجابة", lines=6)
    
    upload_btn.click(fn=upload_file, inputs=pdf_input, outputs=upload_output)
    question_input.submit(fn=answer_question, inputs=question_input, outputs=answer_output)

# تشغيل Gradio
if __name__ == "__main__":
    demo.launch()