"""
RAG Engine — يقرأ ملفات PDF/TXT ويبني vector store محلي
ثم يجيب على الأسئلة باستخدام Google Gemini
"""

import os
import pickle
from pathlib import Path
from typing import List, Tuple

import google.generativeai as genai
import numpy as np

# ────────────────────────────────────────────────
# 1. إعداد Gemini
# ────────────────────────────────────────────────
def init_gemini(api_key: str):
    genai.configure(api_key=api_key)


# ────────────────────────────────────────────────
# 2. قراءة الملفات وتقطيعها (chunking)
# ────────────────────────────────────────────────
def load_text_file(path: str) -> str:
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        return f.read()


def load_pdf_file(path: str) -> str:
    try:
        import pypdf
        reader = pypdf.PdfReader(path)
        return "\n".join(page.extract_text() or "" for page in reader.pages)
    except ImportError:
        raise ImportError("مكتبة pypdf مش متاحة. ضيف pypdf في requirements.txt")


def load_document(path: str) -> str:
    ext = Path(path).suffix.lower()
    if ext == ".pdf":
        return load_pdf_file(path)
    elif ext in (".txt", ".md", ".csv"):
        return load_text_file(path)
    else:
        raise ValueError(f"نوع الملف غير مدعوم: {ext}")


def chunk_text(text: str, chunk_size: int = 800, overlap: int = 100) -> List[str]:
    """يقسّم النص لأجزاء متداخلة."""
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunk = " ".join(words[i : i + chunk_size])
        chunks.append(chunk)
        i += chunk_size - overlap
    return [c for c in chunks if c.strip()]


# ────────────────────────────────────────────────
# 3. Embedding باستخدام Gemini
# ────────────────────────────────────────────────
def embed_texts(texts: List[str], task_type: str = "retrieval_document") -> List[List[float]]:
    """يحوّل النصوص لـ vectors باستخدام Gemini embedding."""
    embeddings = []
    for text in texts:
        result = genai.embed_content(
            model="models/text-embedding-004",
            content=text,
            task_type=task_type,
        )
        embeddings.append(result["embedding"])
    return embeddings


def cosine_similarity(a: List[float], b: List[float]) -> float:
    a, b = np.array(a), np.array(b)
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))


# ────────────────────────────────────────────────
# 4. Vector Store (محلي بدون قاعدة بيانات خارجية)
# ────────────────────────────────────────────────
STORE_PATH = "docs/vector_store.pkl"


class SimpleVectorStore:
    def __init__(self):
        self.chunks: List[str] = []
        self.embeddings: List[List[float]] = []
        self.sources: List[str] = []

    def add_documents(self, file_path: str):
        text = load_document(file_path)
        chunks = chunk_text(text)
        vecs = embed_texts(chunks)
        self.chunks.extend(chunks)
        self.embeddings.extend(vecs)
        self.sources.extend([Path(file_path).name] * len(chunks))
        print(f"✅ تمت إضافة {len(chunks)} chunk من {Path(file_path).name}")

    def search(self, query: str, top_k: int = 4) -> List[Tuple[str, str, float]]:
        if not self.embeddings:
            return []
        q_vec = embed_texts([query], task_type="retrieval_query")[0]
        scores = [cosine_similarity(q_vec, e) for e in self.embeddings]
        top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:top_k]
        return [(self.chunks[i], self.sources[i], scores[i]) for i in top_indices]

    def save(self, path: str = STORE_PATH):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "wb") as f:
            pickle.dump(self, f)

    @staticmethod
    def load(path: str = STORE_PATH) -> "SimpleVectorStore":
        with open(path, "rb") as f:
            return pickle.load(f)

    @staticmethod
    def exists(path: str = STORE_PATH) -> bool:
        return os.path.exists(path)


# ────────────────────────────────────────────────
# 5. الـ Chatbot مع RAG
# ────────────────────────────────────────────────
class GeminiRAGAgent:
    def __init__(self, api_key: str, store: SimpleVectorStore):
        init_gemini(api_key)
        self.store = store
        self.model = genai.GenerativeModel("gemini-2.0-flash")
        self.chat = self.model.start_chat(history=[])

    def answer(self, question: str, use_rag: bool = True) -> Tuple[str, List[str]]:
        context_parts = []
        sources = []

        if use_rag and self.store.chunks:
            results = self.store.search(question, top_k=4)
            for chunk, src, score in results:
                if score > 0.4:          # فلتر الـ chunks الضعيفة
                    context_parts.append(chunk)
                    if src not in sources:
                        sources.append(src)

        if context_parts:
            context = "\n\n---\n\n".join(context_parts)
            prompt = (
                f"أنت مساعد ذكي. استخدم السياق التالي للإجابة على السؤال.\n"
                f"إذا السياق مش بيساعد، قل ذلك وجاوب من معلوماتك.\n\n"
                f"السياق:\n{context}\n\n"
                f"السؤال: {question}"
            )
        else:
            prompt = question

        response = self.chat.send_message(prompt)
        return response.text, sources

"""
Streamlit UI — واجهة الـ AI Agent مع RAG
"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))

import streamlit as st
from rag_engine import GeminiRAGAgent, SimpleVectorStore, STORE_PATH

# ─── إعدادات الصفحة ───────────────────────────
st.set_page_config(
    page_title="🤖 AI Agent — Gemini RAG",
    page_icon="🤖",
    layout="wide",
)

st.title("🤖 AI Agent بـ Google Gemini + RAG")
st.caption("ارفع ملفاتك وابدأ تسأل!")

# ─── Sidebar: API Key + رفع ملفات ─────────────
with st.sidebar:
    st.header("⚙️ الإعدادات")
    api_key = st.text_input("🔑 Gemini API Key", type="password", placeholder="AIza...")

    st.divider()
    st.header("📂 رفع ملفات (RAG)")
    uploaded_files = st.file_uploader(
        "ارفع PDF أو TXT",
        type=["pdf", "txt", "md"],
        accept_multiple_files=True,
    )

    if uploaded_files and api_key:
        if st.button("📥 إضافة الملفات للـ Agent", use_container_width=True):
            # تحميل أو إنشاء vector store
            store = SimpleVectorStore.load(STORE_PATH) if SimpleVectorStore.exists(STORE_PATH) else SimpleVectorStore()

            import tempfile, pathlib
            with st.spinner("جاري معالجة الملفات..."):
                for uf in uploaded_files:
                    suffix = pathlib.Path(uf.name).suffix
                    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
                        tmp.write(uf.getvalue())
                        tmp_path = tmp.name
                    try:
                        import google.generativeai as genai
                        genai.configure(api_key=api_key)
                        store.add_documents(tmp_path)
                        st.success(f"✅ تمت إضافة: {uf.name}")
                    except Exception as e:
                        st.error(f"❌ خطأ في {uf.name}: {e}")
                    finally:
                        os.unlink(tmp_path)
                store.save(STORE_PATH)
            st.success("🎉 تم حفظ الـ vector store!")

    st.divider()
    st.header("🗄️ حالة قاعدة البيانات")
    if SimpleVectorStore.exists(STORE_PATH):
        try:
            _store = SimpleVectorStore.load(STORE_PATH)
            st.info(f"📊 {len(_store.chunks)} chunk محفوظ\n\nالمصادر: {list(set(_store.sources))}")
        except Exception:
            st.warning("⚠️ Vector store موجود لكن تعذّر قراءته")
    else:
        st.info("لا يوجد ملفات مرفوعة بعد")

    if st.button("🗑️ مسح كل الملفات", use_container_width=True):
        if os.path.exists(STORE_PATH):
            os.remove(STORE_PATH)
        st.rerun()

# ─── تهيئة session state ───────────────────────
if "messages" not in st.session_state:
    st.session_state.messages = []

if "agent" not in st.session_state:
    st.session_state.agent = None

# تهيئة الـ agent إذا في API key
if api_key:
    try:
        store = SimpleVectorStore.load(STORE_PATH) if SimpleVectorStore.exists(STORE_PATH) else SimpleVectorStore()
        if st.session_state.agent is None:
            st.session_state.agent = GeminiRAGAgent(api_key=api_key, store=store)
    except Exception as e:
        st.error(f"خطأ في تهيئة الـ agent: {e}")

# ─── عرض المحادثة ──────────────────────────────
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])
        if msg.get("sources"):
            st.caption(f"📎 المصادر: {', '.join(msg['sources'])}")

# ─── مربع الإدخال ──────────────────────────────
if question := st.chat_input("اكتب سؤالك هنا..."):
    if not api_key:
        st.warning("⚠️ ضيف الـ Gemini API Key أولاً من الـ sidebar")
        st.stop()

    # رسالة المستخدم
    st.session_state.messages.append({"role": "user", "content": question})
    with st.chat_message("user"):
        st.markdown(question)

    # رد الـ agent
    with st.chat_message("assistant"):
        with st.spinner("جاري التفكير..."):
            try:
                answer, sources = st.session_state.agent.answer(question)
                st.markdown(answer)
                if sources:
                    st.caption(f"📎 المصادر: {', '.join(sources)}")
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": answer,
                    "sources": sources,
                })
            except Exception as e:
                err = f"❌ حدث خطأ: {e}"
                st.error(err)
                st.session_state.messages.append({"role": "assistant", "content": err})

# ─── زرار مسح المحادثة ────────────────────────
if st.session_state.messages:
    if st.button("🗑️ مسح المحادثة"):
        st.session_state.messages = []
        st.session_state.agent = None
        st.rerun()

streamlit>=1.35.0
google-generativeai>=0.7.0
pypdf>=4.0.0
numpy>=1.26.0

# 🤖 AI Agent — Google Gemini + RAG

Chatbot ذكي مبني على **Google Gemini** يقدر يبحث في ملفاتك (PDF/TXT) ويجاوب على أسئلتك.

---

## 🗂️ هيكل المشروع

```
gemini-rag-agent/
├── app.py                  ← واجهة Streamlit
├── src/
│   └── rag_engine.py       ← الـ RAG engine الأساسي
├── docs/
│   └── vector_store.pkl    ← يتخلق تلقائياً بعد رفع الملفات
├── requirements.txt
└── README.md
```

---

## 🚀 خطوات التشغيل

### 1. احصل على Gemini API Key (مجاناً)
اذهب إلى [Google AI Studio](https://aistudio.google.com/app/apikey) وأنشئ API Key.

### 2. ثبّت المكتبات

```bash
pip install -r requirements.txt
```

### 3. شغّل التطبيق

```bash
streamlit run app.py
```

### 4. استخدمه
1. **ضع الـ API Key** في الـ sidebar
2. **ارفع ملفاتك** (PDF أو TXT) واضغط "إضافة الملفات"
3. **ابدأ تسأل** في مربع الدردشة

---

## ⚙️ كيف يعمل؟

```
ملفاتك (PDF/TXT)
      ↓
   تقطيع (chunking)
      ↓
   Gemini Embedding
      ↓
   Vector Store (محلي)
      ↓
سؤالك → بحث في الـ chunks → أقرب 4 أجزاء
      ↓
   Gemini يجاوب مع السياق
```

---

## 🔧 إعدادات قابلة للتعديل

في `src/rag_engine.py`:

| الإعداد | القيمة الافتراضية | الوصف |
|---------|------------------|-------|
| `chunk_size` | 800 كلمة | حجم كل جزء من النص |
| `overlap` | 100 كلمة | التداخل بين الأجزاء |
| `top_k` | 4 | عدد الأجزاء المسترجعة |
| `score > 0.4` | 0.4 | حد الثقة للسياق |

---

## 📦 المكتبات المستخدمة

- **[google-generativeai](https://pypi.org/project/google-generativeai/)** — Gemini LLM + Embedding
- **[streamlit](https://streamlit.io/)** — واجهة المستخدم
- **[pypdf](https://pypi.org/project/pypdf/)** — قراءة ملفات PDF
- **numpy** — حسابات التشابه (cosine similarity)