"""
نظام RAG للـ HR Chatbot باستخدام Ollama و Llama 3.1
HR Chatbot RAG System with Ollama & Llama 3.1

هذه النسخة تستخدم:
- sentence-transformers للـ Embeddings
- FAISS للـ Vector Store
- Ollama API مباشرة للـ LLM
"""

import os
import json
import requests
import numpy as np
from typing import List, Dict, Optional, Tuple


class Document:
    """مستند بسيط للـ RAG"""
    def __init__(self, page_content: str, metadata: dict = None):
        self.page_content = page_content
        self.metadata = metadata or {}


class SimpleVectorStore:
    """Vector Store بسيط باستخدام FAISS"""
    
    def __init__(self, embedding_model):
        self.embedding_model = embedding_model
        self.documents: List[Document] = []
        self.embeddings: Optional[np.ndarray] = None
        self.index = None
        
    def add_documents(self, documents: List[Document]):
        """إضافة مستندات"""
        # استخراج النصوص
        texts = [doc.page_content for doc in documents]
        
        # إنشاء embeddings
        new_embeddings = self.embedding_model.encode(texts, convert_to_numpy=True)
        
        if self.embeddings is None:
            self.embeddings = new_embeddings
            self.documents = documents
        else:
            self.embeddings = np.vstack([self.embeddings, new_embeddings])
            self.documents.extend(documents)
        
        # بناء FAISS index
        self._build_index()
        
    def _build_index(self):
        """بناء FAISS index"""
        import faiss
        
        if self.embeddings is None:
            return
            
        # تطبيع الـ embeddings
        embeddings_copy = self.embeddings.copy().astype('float32')
        faiss.normalize_L2(embeddings_copy)
        
        # إنشاء index
        dimension = embeddings_copy.shape[1]
        self.index = faiss.IndexFlatIP(dimension)  # Inner Product (Cosine Similarity)
        self.index.add(embeddings_copy)
        
    def similarity_search(self, query: str, k: int = 5) -> List[Document]:
        """بحث التشابه"""
        if self.index is None:
            return []
        
        import faiss
        
        # تحويل السؤال إلى embedding
        query_embedding = self.embedding_model.encode([query], convert_to_numpy=True).astype('float32')
        faiss.normalize_L2(query_embedding)
        
        # البحث
        scores, indices = self.index.search(query_embedding, k)
        
        results = []
        for idx in indices[0]:
            if 0 <= idx < len(self.documents):
                results.append(self.documents[idx])
        
        return results
    
    def similarity_search_with_score(self, query: str, k: int = 5) -> List[Tuple[Document, float]]:
        """بحث التشابه مع الدرجات"""
        if self.index is None:
            return []
        
        import faiss
        
        # تحويل السؤال إلى embedding
        query_embedding = self.embedding_model.encode([query], convert_to_numpy=True).astype('float32')
        faiss.normalize_L2(query_embedding)
        
        # البحث
        scores, indices = self.index.search(query_embedding, k)
        
        results = []
        for score, idx in zip(scores[0], indices[0]):
            if 0 <= idx < len(self.documents):
                results.append((self.documents[idx], float(score)))
        
        return results
    
    def save(self, path: str):
        """حفظ Vector Store"""
        import faiss
        
        os.makedirs(path, exist_ok=True)
        
        # حفظ FAISS index
        if self.index is not None:
            faiss.write_index(self.index, os.path.join(path, "index.faiss"))
        
        # حفظ المستندات
        docs_data = [{"content": doc.page_content, "metadata": doc.metadata} for doc in self.documents]
        with open(os.path.join(path, "documents.json"), "w", encoding="utf-8") as f:
            json.dump(docs_data, f, ensure_ascii=False, indent=2)
        
        # حفظ embeddings
        if self.embeddings is not None:
            np.save(os.path.join(path, "embeddings.npy"), self.embeddings)
    
    def load(self, path: str) -> bool:
        """تحميل Vector Store"""
        import faiss
        
        try:
            # تحميل FAISS index
            index_path = os.path.join(path, "index.faiss")
            if os.path.exists(index_path):
                self.index = faiss.read_index(index_path)
            
            # تحميل المستندات
            docs_path = os.path.join(path, "documents.json")
            if os.path.exists(docs_path):
                with open(docs_path, "r", encoding="utf-8") as f:
                    docs_data = json.load(f)
                self.documents = [Document(d["content"], d["metadata"]) for d in docs_data]
            
            # تحميل embeddings
            emb_path = os.path.join(path, "embeddings.npy")
            if os.path.exists(emb_path):
                self.embeddings = np.load(emb_path)
            
            return True
        except Exception as e:
            print(f"❌ خطأ في تحميل Vector Store: {e}")
            return False


class HRChatbotRAG:
    """نظام RAG للـ HR Chatbot باستخدام Ollama"""
    
    def __init__(
        self,
        embedding_model: str = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
        llm_model: str = "llama3.1",
        ollama_base_url: str = "http://localhost:11434",
        temperature: float = 0.3,
        streaming: bool = False
    ):
        """
        تهيئة نظام RAG
        
        Args:
            embedding_model: نموذج الـ Embeddings
            llm_model: نموذج Ollama (llama3.1)
            ollama_base_url: رابط Ollama
            temperature: درجة الإبداعية
            streaming: تفعيل البث المباشر
        """
        self.embedding_model_name = embedding_model
        self.llm_model = llm_model
        self.ollama_base_url = ollama_base_url
        self.temperature = temperature
        self.streaming = streaming
        
        print(f"🚀 جاري تهيئة نظام RAG...")
        print(f"   📦 Embedding Model: {embedding_model}")
        print(f"   🤖 LLM Model: {llm_model} (Ollama)")
        
        # تهيئة Embeddings
        print("   ⏳ تحميل نموذج Embeddings...")
        from sentence_transformers import SentenceTransformer
        self.embedding_model = SentenceTransformer(embedding_model)
        
        # تهيئة Vector Store
        self.vectorstore = SimpleVectorStore(self.embedding_model)
        
        # التحقق من الاتصال بـ Ollama
        print("   ⏳ الاتصال بـ Ollama...")
        if self._test_ollama_connection():
            print(f"   ✅ تم الاتصال بـ Ollama بنجاح!")
        else:
            print(f"   ⚠️ تحذير: تأكد من تشغيل Ollama!")
        
        print("✅ تم تهيئة النظام بنجاح!")
    
    def _test_ollama_connection(self) -> bool:
        """اختبار الاتصال بـ Ollama"""
        try:
            response = requests.get(f"{self.ollama_base_url}/api/tags", timeout=5)
            return response.status_code == 200
        except:
            return False
    
    def _call_ollama(self, prompt: str) -> str:
        """استدعاء Ollama API"""
        try:
            response = requests.post(
                f"{self.ollama_base_url}/api/generate",
                json={
                    "model": self.llm_model,
                    "prompt": prompt,
                    "stream": False,
                    "options": {
                        "temperature": self.temperature,
                        "num_ctx": 4096,
                        "num_predict": 512
                    }
                },
                timeout=120
            )
            
            if response.status_code == 200:
                return response.json().get("response", "")
            else:
                return f"خطأ في Ollama: {response.status_code}"
        except Exception as e:
            return f"خطأ في الاتصال بـ Ollama: {str(e)}"
    
    def _create_prompt(self, context: str, question: str) -> str:
        """إنشاء Prompt للـ HR Chatbot"""
        return f"""أنت مساعد ذكي لنظام الموارد البشرية (HR Genie). مهمتك هي الإجابة على أسئلة الموظفين بشكل دقيق ومهني باللغة العربية.

استخدم المعلومات التالية للإجابة على السؤال. إذا لم تجد الإجابة في المعلومات المتاحة، قل ذلك بوضوح واقترح التواصل مع قسم الموارد البشرية.

المعلومات المتاحة:
{context}

السؤال: {question}

التعليمات:
- أجب باللغة العربية بشكل واضح ومختصر
- إذا كان السؤال عن إجراء معين، اشرح الخطوات بوضوح
- كن ودوداً ومهنياً في الرد
- إذا لم تكن متأكداً من الإجابة، اقترح التواصل مع قسم HR

الإجابة:"""
    
    def build_vectorstore(self, documents: List[Document]) -> None:
        """
        بناء Vector Store من المستندات
        
        Args:
            documents: قائمة المستندات
        """
        print(f"🔨 جاري بناء Vector Store من {len(documents)} مستند...")
        self.vectorstore.add_documents(documents)
        print(f"✅ تم بناء Vector Store بنجاح!")
    
    def save_vectorstore(self, path: str = "embeddings/vectorstore") -> None:
        """حفظ Vector Store"""
        self.vectorstore.save(path)
        print(f"💾 تم حفظ Vector Store في: {path}")
    
    def load_vectorstore(self, path: str = "embeddings/vectorstore") -> bool:
        """تحميل Vector Store محفوظ"""
        if self.vectorstore.load(path):
            print(f"📂 تم تحميل Vector Store من: {path}")
            return True
        return False
    
    def ask(self, question: str, k: int = 5) -> Dict:
        """
        طرح سؤال على الـ Chatbot
        
        Args:
            question: السؤال
            k: عدد المستندات للاسترجاع
        
        Returns:
            dict مع الإجابة والمصادر
        """
        if not self.vectorstore.documents:
            return {
                "success": False,
                "answer": "❌ النظام غير جاهز. يرجى بناء قاعدة المعرفة أولاً.",
                "sources": [],
                "error": "Vector store not initialized"
            }
        
        try:
            # البحث عن المستندات المتشابهة
            docs = self.vectorstore.similarity_search(question, k=k)
            
            # إنشاء السياق
            context = "\n\n".join([doc.page_content for doc in docs])
            
            # إنشاء Prompt والاستدعاء
            prompt = self._create_prompt(context, question)
            answer = self._call_ollama(prompt)
            
            # استخراج المصادر
            sources = []
            for doc in docs:
                sources.append({
                    "content": doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content,
                    "metadata": doc.metadata
                })
            
            return {
                "success": True,
                "answer": answer,
                "sources": sources,
                "question": question
            }
            
        except Exception as e:
            return {
                "success": False,
                "answer": f"❌ حدث خطأ: {str(e)}",
                "sources": [],
                "error": str(e)
            }
    
    def similarity_search(self, query: str, k: int = 5) -> List[Document]:
        """بحث التشابه"""
        return self.vectorstore.similarity_search(query, k)
    
    def similarity_search_with_score(self, query: str, k: int = 5) -> List[Tuple[Document, float]]:
        """بحث التشابه مع الدرجات"""
        return self.vectorstore.similarity_search_with_score(query, k)
    
    def add_documents(self, documents: List[Document]) -> None:
        """إضافة مستندات جديدة"""
        self.vectorstore.add_documents(documents)
        print(f"➕ تم إضافة {len(documents)} مستند جديد")
    
    def get_relevant_context(self, question: str, k: int = 3) -> str:
        """الحصول على السياق المناسب للسؤال"""
        docs = self.similarity_search(question, k=k)
        return "\n\n".join([doc.page_content for doc in docs])
    
    def get_stats(self) -> Dict:
        """الحصول على إحصائيات النظام"""
        return {
            "embedding_model": self.embedding_model_name,
            "llm_model": self.llm_model,
            "ollama_url": self.ollama_base_url,
            "vectorstore_ready": len(self.vectorstore.documents) > 0,
            "vectorstore_size": len(self.vectorstore.documents)
        }


if __name__ == "__main__":
    # اختبار سريع
    print("=" * 50)
    print("🧪 اختبار نظام RAG")
    print("=" * 50)
    
    # تهيئة النظام
    rag = HRChatbotRAG(llm_model="llama3.1")
    
    # إنشاء مستندات تجريبية
    test_docs = [
        Document(
            page_content="السؤال: كم عدد أيام الإجازة السنوية؟\nالإجابة: يحق لكل موظف 21 يوم إجازة سنوية مدفوعة الأجر بعد إكمال سنة من العمل.",
            metadata={"source": "test", "type": "qa"}
        ),
        Document(
            page_content="السؤال: متى يتم صرف الراتب؟\nالإجابة: يتم صرف الرواتب يوم 27 من كل شهر. لو اليوم ده إجازة رسمية، بينزل المرتب في يوم العمل اللي قبله.",
            metadata={"source": "test", "type": "qa"}
        ),
    ]
    
    # بناء Vector Store
    rag.build_vectorstore(test_docs)
    
    # اختبار سؤال
    print("\n" + "=" * 50)
    result = rag.ask("كم يوم إجازة عندي؟")
    print(f"✅ الإجابة: {result['answer']}")
    print("=" * 50)