import os
import uuid
import gradio as gr
import chromadb
from openai import OpenAI
from chromadb.utils import embedding_functions
from pypdf import PdfReader
from docx import Document

# --- 1. إعداد الذاكرة السيبرانية (ChromaDB) ---
CHROMA_DATA_PATH = "./neural_memory_v2"
chroma_client = chromadb.PersistentClient(path=CHROMA_DATA_PATH)
ef = embedding_functions.DefaultEmbeddingFunction()
collection = chroma_client.get_or_create_collection(name="yousef_vault", embedding_function=ef)

# النموذج الأساسي الافتراضي (يمكنك تغييره من هنا)
DEFAULT_MODEL = "huihui-ai/Qwen2.5-72B-Instruct-abliterated"

# --- 2. دالة استخراج النصوص (يدعم PDF, DOCX, TXT) ---
def process_document(file):
    if file is None: return "⚠️ لم يتم اختيار ملف."
    text = ""
    try:
        file_ext = os.path.splitext(file.name)[1].lower()
        if file_ext == ".pdf":
            reader = PdfReader(file.name)
            for page in reader.pages:
                text += (page.extract_text() or "") + "\n"
        elif file_ext == ".docx":
            doc = Document(file.name)
            text = "\n".join([para.text for para in doc.paragraphs])
        elif file_ext in [".txt", ".md"]:
            with open(file.name, "r", encoding="utf-8", errors="ignore") as f:
                text = f.read()
        
        if not text.strip():
            return "❌ تعذر استخراج نص. تأكد أن الملف ليس صوراً فقط."

        # تقسيم النص لقطع وحقنها في الذاكرة
        chunks = [text[i:i+2000] for i in range(0, len(text), 2000)]
        ids = [str(uuid.uuid4()) for _ in chunks]
        collection.add(documents=chunks, ids=ids)
        
        return f"✅ تمت الأرشفة: {len(chunks)} شظية معرفية حُقنت في الذاكرة."
    except Exception as e:
        return f"❌ خطأ في المعالجة: {str(e)}"

# --- 3. دالة الحوار (RAG) مع معالجة أخطاء الـ Stream ---
def predict(message, history, system_prompt, temperature, custom_model, oauth_token: gr.OAuthToken | None):
    # استخدام توكن المستخدم المسجل (OAuth) لخصم الرصيد الشخصي
    api_key = oauth_token.token if oauth_token else os.getenv("HF_TOKEN")
    
    if not api_key:
        yield "⚠️ يرجى تسجيل الدخول عبر Hugging Face أولاً للوصول إلى رصيدك الشخصي."
        return

    client = OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=api_key
    )
    
    # تحديد النموذج المستخدم
    active_model = custom_model.strip() if custom_model and custom_model.strip() else DEFAULT_MODEL
    
    # استرجاع سياق من الذاكرة
    results = collection.query(query_texts=[message], n_results=10)
    context = "\n".join(results['documents'][0]) if results['documents'] else ""
    
    # بناء قائمة الرسائل
    messages = [{"role": "system", "content": f"{system_prompt}\n\n[CONTEXT]:\n{context}"}]
    
    for msg in history:
        messages.append({"role": msg['role'], "content": msg['content']})
        
    messages.append({"role": "user", "content": message})

    try:
        response = client.chat.completions.create(
            model=active_model,
            messages=messages,
            temperature=temperature,
            stream=True
        )

        partial_message = ""
        for chunk in response:
            # حماية ضد خطأ list index out of range
            if hasattr(chunk, 'choices') and chunk.choices and len(chunk.choices) > 0:
                content = chunk.choices[0].delta.content
                if content:
                    partial_message += content
                    yield partial_message
            else:
                # تجاهل الإطارات الفارغة التي تظهر في نهاية البث
                continue
                
    except Exception as e:
        if partial_message:
            # إذا كان هناك نص تم توليده بالفعل، اعرضه ولا تظهر رسالة الخطأ
            yield partial_message
        else:
            yield f"❌ فشل الاتصال بالنموذج ({active_model}): {str(e)}"

# --- 4. بناء الواجهة (بسيطة، داكنة، ومتجاوبة) ---
with gr.Blocks(fill_height=True, theme=gr.themes.Soft(primary_hue="indigo")) as demo:
    with gr.Sidebar():
        gr.Markdown("# 🧬 دهليز يوسف")
        gr.LoginButton("Sign in")
        
        with gr.Accordion("⚙️ الضبط والوعي", open=True):
            model_input = gr.Textbox(
                label="المحرك المستهدف (Model ID)",
                placeholder="اتركه فارغاً للنموذج الأساسي",
                value="",
                lines=1
            )
            system_input = gr.Textbox(
                value="أنت رفيق حكيم يغوص في التيه مع المستخدم. أجب بصدق مستعيناً بالذاكرة.", 
                label="البرومبت النظامي"
            )
            temp_slider = gr.Slider(0.1, 1.5, 0.8, label="درجة التشظي (Temperature)")
        
        gr.Markdown("---")
        gr.Markdown("### 📂 ملقم الذاكرة (PDF, DOCX, TXT)")
        file_box = gr.File(label="ارفع مسوداتك هنا", file_types=[".pdf", ".docx", ".txt"])
        status_msg = gr.Markdown("*الذاكرة بانتظار الملفات...*")
        
        file_box.change(process_document, inputs=file_box, outputs=status_msg)
        
        clear_btn = gr.Button("🗑 مسح الذاكرة")
        def clear_memory():
            all_data = collection.get()
            if all_data['ids']: 
                collection.delete(ids=all_data['ids'])
            return "🗑 تم تصفير الذاكرة بالكامل."
        clear_btn.click(clear_memory, outputs=status_msg)

    gr.ChatInterface(
        predict,
        additional_inputs=[system_input, temp_slider, model_input],
        fill_height=True,
        type="messages"
    )

if __name__ == "__main__":
    demo.launch()