import dspy
import gradio as gr
import chromadb
import fitz  # PyMuPDF
import json
from sentence_transformers import SentenceTransformer
from dspy import Example, MIPROv2, evaluate

# إعداد النموذج المفتوح المصدر
dspy.settings.configure(lm=dspy.LM("mistralai/Mistral-7B-Instruct-v0.2"))

# إعداد Chroma
client = chromadb.PersistentClient(path="./chroma_db")
col = client.get_or_create_collection(name="arabic_docs")

# إعداد نموذج التضمين (Embeddings)
embedder = SentenceTransformer("sentence-transformers/LaBSE")

# تقطيع نصوص PDF
def process_pdf(pdf_bytes):
    doc = fitz.open(stream=pdf_bytes, filetype="pdf")
    texts = []
    for p in doc:
        text = p.get_text()
        for chunk in text.split("\n\n"):
            if len(chunk.strip()) > 50:
                texts.append(chunk.strip())
    return texts

# إدخال البيانات إلى Chroma مع دعم NamedString
def ingest(pdf_file):
    if hasattr(pdf_file, "read"):
        pdf_bytes = pdf_file.read()
    else:
        pdf_bytes = pdf_file  # قد يكون bytes أو NamedString
    texts = process_pdf(pdf_bytes)
    embeddings = embedder.encode(texts, show_progress_bar=True)
    for i, (chunk, emb) in enumerate(zip(texts, embeddings)):
        col.add(ids=[f"chunk_{i}"], embeddings=[emb.tolist()], metadatas=[{"text": chunk}])
    return f"✅ تمت إضافة {len(texts)} مقطعاً."

# استرجاع السياق من Chroma
def retrieve_context(query):
    query_emb = embedder.encode([query])[0]
    results = col.query(query_embeddings=[query_emb.tolist()], n_results=1)
    context_list = [m["text"] for group in results["metadatas"] for m in group]
    return context_list[0] if context_list else ""

# تعريف توقيع DSPy
class RagSig(dspy.Signature):
    question: str = dspy.InputField()
    context: str = dspy.InputField()
    answer: str = dspy.OutputField()

# وحدة DSPy
class RagMod(dspy.Module):
    def __init__(self):
        super().__init__()
        self.predictor = dspy.Predict(RagSig)

    def forward(self, question):
        context = retrieve_context(question)
        return self.predictor(question=question, context=context)

model = RagMod()

# توليد إجابة
def answer(question):
    out = model(question)
    return out.answer

# تحميل بيانات التدريب/التقييم
def load_dataset(path):
    with open(path, "r", encoding="utf-8") as f:
        return [Example(**json.loads(l)).with_inputs("question") for l in f]

# تحسين النموذج
def optimize(train_file, val_file):
    global model
    trainset = load_dataset(train_file.name)
    valset = load_dataset(val_file.name)
    tp = MIPROv2(metric=evaluate.answer_exact_match, auto="light", num_threads=4)
    optimized = tp.compile(model, trainset=trainset, valset=valset)
    model = optimized
    return "✅ تم تحسين النموذج!"

# واجهة Gradio
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 نظام RAG عربي باستخدام DSPy + نموذج مفتوح المصدر")

    with gr.Tab("📥 تحميل وتخزين"):
        pdf_input = gr.File(label="ارفع ملف PDF", file_types=[".pdf"])
        ingest_btn = gr.Button("إضافة إلى قاعدة البيانات")
        ingest_output = gr.Textbox()
        ingest_btn.click(fn=ingest, inputs=pdf_input, outputs=ingest_output)

    with gr.Tab("❓ سؤال"):
        q = gr.Textbox(label="اكتب سؤالك بالعربية")
        answer_btn = gr.Button("احصل على الإجابة")
        out = gr.Textbox(label="الإجابة")
        answer_btn.click(fn=answer, inputs=q, outputs=out)

    with gr.Tab("⚙️ تحسين النموذج"):
        train_file = gr.File(label="trainset.jsonl", file_types=[".jsonl"])
        val_file = gr.File(label="valset.jsonl", file_types=[".jsonl"])
        opt_btn = gr.Button("ابدأ التحسين")
        result = gr.Textbox(label="نتيجة التحسين")
        opt_btn.click(optimize, inputs=[train_file, val_file], outputs=result)

    demo.launch()