#import libraries
import os
from datetime import datetime
import gradio as gr
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA, LLMChain
from langchain.prompts import PromptTemplate
import json
import os
import json
import time
from datetime import datetime
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain.llms import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
#import libraries
import os
from datetime import datetime
import gradio as gr
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_core.runnables import RunnablePassthrough
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA, LLMChain
from langchain.prompts import PromptTemplate
import json

# tracking answer
LOG_FILE = "chat_log.txt"
TRACKING_FILE = "source_tracking.json"
ANSWERS_FILE = "answers_tracking.txt"

def init_files():
    if not os.path.exists(LOG_FILE):
        with open(LOG_FILE, 'w', encoding='utf-8') as f:
            f.write("سجل محادثات التأمين الصحي\n")
            f.write("="*50 + "\n\n")

    if not os.path.exists(TRACKING_FILE):
        with open(TRACKING_FILE, 'w', encoding='utf-8') as f:
            json.dump([], f)

    if not os.path.exists(ANSWERS_FILE):
        with open(ANSWERS_FILE, 'w', encoding='utf-8') as f:
            f.write("سجل الإجابات والمصادر\n")
            f.write("="*50 + "\n\n")

def write_to_log(*messages):
    with open(LOG_FILE, 'a', encoding='utf-8') as f:
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        for message in messages:
            f.write(f"{timestamp} - {message}\n")

def track_sources(question, answer, sources):
    # 1.    JSON
    with open(TRACKING_FILE, 'r+', encoding='utf-8') as f:
        data = json.load(f)
        entry = {
            "timestamp": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            "question": question,
            "answer": answer,
            "sources": [
                {
                    "file": doc.metadata['source'],
                    "section": doc.metadata.get('section', 'غير محدد'),
                    "content": doc.page_content
                }
                for doc in sources
            ]
        }
        data.append(entry)
        f.seek(0)
        json.dump(data, f, ensure_ascii=False, indent=2)

    # 2. ANSWERS_FILE
    with open(ANSWERS_FILE, 'a', encoding='utf-8') as f:
        f.write(f"\n{'='*100}\n")
        f.write(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]\n")
        f.write(f"السؤال الأصلي: {question}\n\n")

        f.write("الإجابة الكاملة:\n")
        f.write(f"{answer}\n\n")

        f.write(f"المصادر المستخدمة ({len(sources)} مصدر):\n")
        for i, doc in enumerate(sources, 1):
            f.write(f"\nالمصدر #{i}:\n")
            f.write(f"- الملف: {os.path.basename(doc.metadata['source'])}\n")
            f.write(f"- العنوان: {doc.metadata.get('section', 'غير محدد')}\n")
            f.write(f"- المحتوى الكامل:\n{doc.page_content}\n")
            f.write("-"*80 + "\n")

        f.write(f"{'='*100}\n")
import os
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain.chains import RetrievalQA
from langchain_openai import ChatOpenAI





def load_embeddings() -> FAISS:
    """Load FAISS index and metadata"""
    required_files = ["index.faiss", "index.pkl", "source_files.txt"]

    missing_files = [file for file in required_files if not os.path.exists(file)]
    if missing_files:
        error_msg = f"الملفات التالية غير موجودة: {', '.join(missing_files)}"
        write_to_log(error_msg)
        raise ValueError(error_msg)

    try:
        embedding_model = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-large")
        vectorstore = FAISS.load_local(
            folder_path=".",
            embeddings=embedding_model,
            allow_dangerous_deserialization=True
        )
        return vectorstore
    except Exception as e:
        raise RuntimeError(f"فشل في تحميل ملفات FAISS: {str(e)}")


def setup_chains(vectorstore: FAISS):
    """Set up rephrasing and QA chains using GPT-OSS-120B"""

    # ✅ Using Groq's endpoint and GPT-OSS-120B model
    
    llm = ChatOpenAI(
    model="openai/gpt-oss-120b",
    base_url="https://api.groq.com/openai/v1",
    api_key=os.getenv("GROQ_API_KEY"),  
    temperature=0.4,
    max_tokens=3000
)

    rephrase_prompt = PromptTemplate.from_template("""
    قم بتحويل العبارة التالية من العامية المصرية إلى اللغة العربية الفصحى مع الالتزام بالتالي:
    الكلمات التاليه  لا تفيرها (زوجه-اجنبي- المولود - الرعايه).
    1. إذا كانت العبارة بالفصحى بالفعل، اتركها كما هي دون تغيير
    2. لا تقم بإضافة أي كلمات أو تعليقات إضافية
    3. حافظ على نفس المعنى بدقة
    4. غير فقط الكلمات العامية إلى فصحى مع الحفاظ على الكلمات الفصيحة كما هي
    السؤال: "{question}"
    السؤال بالفصحى:
    """)

    rephrase_chain = (
        {"question": RunnablePassthrough()}
        | rephrase_prompt
        | llm
    )

    qa_prompt = PromptTemplate.from_template("""
    أجب على السؤال التالي بناءً على المعلومات الموجودة في النصوص المقدمة لك.

    اشتراطات الإجابة:
    1. اذا كان السوال به اكتر من جزء فتجيب عن كل جزء فالسوال ولا تترك شي
    2. الإجابة يجب أن تكون كاملة دون حذف شي من النص
    3. ذكر جميع الاوراق أو المستندات المطلوبة
    4. التزم باللغة العربية فقط
    5. اذا كان المطلوب اوراق طفل مولود فلا تذكر بطاقه الرقم القومي للمولود بدلا منها اذكر شهاده الميلاد
    6. التامين الصحي شامل هذه المناطق فقط (بورسعيد، الإسماعيلية، السويس، جنوب سيناء، الأقصر، وأسوان)
    7. لا تذكر أرقام خطوات أو إجراءات
    8. لا تشير إلى مصدر المعلومة
    9. لا تضيف أي معلومات خارجية
    10. عند السؤال عن المستندات المطلوبة، يجب ذكر جميع أنواع المستندات (التسجيل، الدخل )

    مثال
      السوال : ما الاوراق المطلوبه لتسجيل طفل مولود

      :الاجابه
        (مستندات التسجيل )
        ***********************************************************
        -يتم استالم صور مع ضرورة االطالع على الاصل
        -صورة بطاقة الرقم القومي )سارية – وجهين(
        -صورة شهادة ميالد مميكنة للطفل
        -أصل بطاقة التأمين الصحي الشامل لألسرة
        -أصل قيد عائلي مميكن إن تطلب األمر

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

    مثال 2
        لو السؤال عن الاجراءات

         يتم ذكر تريتيب الاجراءات  حسب الاحداث المذكوره فالنصوص ثم  يذكر المستندات المطلوبه من مستندات تسجيل ومستندات الدخل

    السؤال: {question}
    النصوص: {context}
    الإجابة:
    """)

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=vectorstore.as_retriever(
            search_type="mmr",
            search_kwargs={
                'k': 12,
                'fetch_k': 100,
                'lambda_mult': 0.6
            }
        ),
        return_source_documents=True,
        chain_type_kwargs={"prompt": qa_prompt}
    )

    return rephrase_chain, qa_chain
    

def display_sources(sources):
    """Display sources with preview"""


def get_answer(question, rephrase_chain, qa_chain):
    """Get answer for a single question"""
    try:
        print("\n" + "="*50)
        print("📝 معالجة السؤال...")
        rewritten = rephrase_chain.invoke({"question": question})
        fusha_question = rewritten.content.strip()
        print(f"السؤال بالفصحى: {fusha_question}")

        for attempt in range(3):
            try:
                # print("🔍 البحث عن الإجابة...")
                result = qa_chain.invoke(fusha_question)
                answer = result["result"]
                sources = result.get("source_documents", [])
                break
            except APIStatusError:
                if attempt == 2: raise
                time.sleep(2 ** (attempt + 5))

        # print("\n💡 الإجابة:")
        # print(answer)
        if sources: display_sources(sources)

        track_sources(question, answer, sources)
        write_to_log(f"سؤال: {question}", f"إجابة: {answer[:200]}...")

        return answer, sources

    except Exception as e:
        error_msg = f"❌ خطأ: {str(e)}"
        print(error_msg)
        raise

def main(question):
    """Main function to get answer for a question"""
    init_files()

    try:
        print("⏳ جار تحميل النموذج...")
        vectorstore = load_embeddings()
        rephrase_chain, qa_chain = setup_chains(vectorstore)

        start_time = time.time()
        answer, sources = get_answer(question, rephrase_chain, qa_chain)
        print(f"\n⏱️ وقت الاستجابة: {time.time() - start_time:.2f} ثانية")

        return {
            "question": question,
            "answer": answer,

        }

    except Exception as e:
        print(f"❌ فشل: {str(e)}")
        return {"error": str(e)}

# Example usage:
if __name__ == "__main__":
    # Ask your question here
    question = "من الجهة المسؤولة عن تسجيل المؤمن عليهم وتحديث بياناتهم؟"
    result = main(question)
    print("\nالنتيجة النهائية:")
    print(json.dumps(result, ensure_ascii=False, indent=2))

# interface

def create_gradio_interface(rephrase_chain, qa_chain):
    with gr.Blocks(title="المساعد الذكي للتأمين الصحي") as demo:
        gr.Markdown("## 🏥 المساعد الذكي للتأمين الصحي")
        gr.Markdown("اسأل عن أي معلومات في وثائق وسياسة التأمين")

        with gr.Row():
            with gr.Column():
                chatbot = gr.Chatbot(label="المحادثة", height=400)
                question_box = gr.Textbox(label="اكتب سؤالك هنا", placeholder="مثال: ما هي مستندات التسجيل  ؟")
                submit_btn = gr.Button("إرسال", variant="primary")

            with gr.Column():
                sources_output = gr.HTML(
                    label="تفاصيل الإجابة والمصادر",
                    value="<div style='border:1px solid #e0e0e0; padding:10px; border-radius:5px; height:400px; overflow-y:auto;'>"
                          "<h4 style='margin-top:0;'>سيظهر هنا تفاصيل المصادر المستخدمة</h4></div>"
                )

        clear_btn = gr.Button("مسح المحادثة")
        chat_history = gr.State([])

        def process_and_display(question, history):
            rewritten = rephrase_chain.invoke({"question": question})
            fusha_question = rewritten.content.strip()

            result = qa_chain.invoke(fusha_question)
            answer = result["result"]


            write_to_log(
                f"السؤال: {question}",
                f"السؤال المحول: {fusha_question}",
                f"الإجابة: {answer}",
                "المصادر المستخدمة:"
            )

            for i, doc in enumerate(result["source_documents"], 1):
                write_to_log(
                    f"المصدر #{i}:",
                    f"الملف: {os.path.basename(doc.metadata['source'])}",
                    f"القسم: {doc.metadata.get('section', 'غير محدد')}",
                    f"المحتوى: {doc.page_content[:300]}..." if len(doc.page_content) > 300 else f"المحتوى: {doc.page_content}",
                    "-"*50
                )

            track_sources(question, answer, result["source_documents"])

            sources_html = """
            <div style='border:1px solid #e0e0e0; padding:10px; border-radius:5px; height:400px; overflow-y:auto;'>
                <h3 style='margin-top:0;'>المصادر المستخدمة:</h3>
                <div style='margin-bottom:20px;'>
                    <h4 style='margin-bottom:5px;'>الإجابة:</h4>
                    <p style='background:#2d3748; padding:10px; border-radius:5px;'>{answer}</p>
                </div>
            """.format(answer=answer)

            for i, doc in enumerate(result["source_documents"][:3], 1):
                sources_html += f"""
                <div style='margin-bottom:15px; border-bottom:1px dashed #e0e0e0; padding-bottom:10px;'>
                    <h4 style='margin-bottom:5px;'>المصدر #{i}</h4>
                    <p><strong>الملف:</strong> {os.path.basename(doc.metadata['source'])}</p>
                    <p><strong>المحتوى:</strong></p>
                    <div style='background:#2d3748; padding:8px; border-radius:3px; max-height:150px; overflow-y:auto;'>
                        {doc.page_content[:500]}{'...' if len(doc.page_content) > 500 else ''}
                    </div>
                </div>
                """

            sources_html += "</div>"

            history.append((question, answer))
            return "", history, sources_html

        submit_btn.click(
            fn=process_and_display,
            inputs=[question_box, chat_history],
            outputs=[question_box, chatbot, sources_output]
        )

        question_box.submit(
            fn=process_and_display,
            inputs=[question_box, chat_history],
            outputs=[question_box, chatbot, sources_output]
        )

        clear_btn.click(
            fn=lambda: ([], "<div style='border:1px solid #e0e0e0; padding:10px; border-radius:5px; height:400px; overflow-y:auto;'>"
                           "<h4 style='margin-top:0;'>سيظهر هنا تفاصيل المصادر المستخدمة</h4></div>"),
            outputs=[chatbot, sources_output],
            inputs=[]
        )

    return demo

if __name__ == "__main__":
    init_files()
    vectorstore = load_embeddings()
    rephrase_chain, qa_chain = setup_chains(vectorstore)
    demo = create_gradio_interface(rephrase_chain, qa_chain)
    demo.launch(server_name="0.0.0.0", server_port=7860)