import os
import logging
from loguru import logger
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.chat_models import ChatOpenAI
from langchain.retrievers import ContextualCompressionRetriever
from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain.prompts import PromptTemplate

# تنظیمات لاگ‌گیری
logger.add("debug.log", rotation="10 MB", level="DEBUG")

# تنظیمات مسیرها
DATABASE_FILE = "/home/user/app/vector_database"  # بدون پسوند .pkl

# بررسی و بارگذاری دیتابیس برداری
def load_database():
    try:
        if os.path.exists(DATABASE_FILE):
            embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("My_huggingface_key"), model="text-embedding-3-large")  # بهبود امبدینگ
            vector_db = FAISS.load_local(DATABASE_FILE, embeddings)
            logger.info(f"✅ دیتابیس بارگذاری شد: {DATABASE_FILE}")
            return vector_db
        else:
            logger.warning(f"❌ دیتابیس در مسیر {DATABASE_FILE} وجود ندارد.")
            return None
    except Exception as e:
        logger.error(f"❌ خطا در بارگذاری دیتابیس: {e}")
        return None

global_vector_db = load_database()

# بهبود بازیابی اسناد با Reranking
def create_retriever(vector_db):
    embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("My_huggingface_key"), model="text-embedding-3-large")
    compressor = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.7)  # تنظیم آستانه تشابه
    retriever = ContextualCompressionRetriever(base_compressor=compressor, base_retriever=vector_db.as_retriever(search_kwargs={"k": 10}))
    return retriever

# پاسخ‌گویی بر اساس اسناد بارگذاری‌شده
def chat_with_doc(query):
    try:
        global global_vector_db
        if not global_vector_db:
            return "❌ دیتابیس بارگذاری نشده است. لطفاً از پشتیبانی کمک بگیرید."
        if not query.strip():
            return "❌ لطفاً سوال خود را وارد کنید."

        # بهبود بازیابی اسناد
        retriever = create_retriever(global_vector_db)
        docs = retriever.get_relevant_documents(query)

        # لاگ‌گیری برای بررسی اسناد بازیابی شده
        logger.info(f"تعداد اسناد بازیابی شده: {len(docs)}")
        for doc in docs:
            logger.info(f"سند بازیابی شده: {doc.page_content[:100]}...")  # نمایش بخشی از متن سند

        context = "\n\n".join([doc.page_content for doc in docs])
        if not context:
            return "هیچ اطلاعات مرتبطی یافت نشد."

        # بهبود Prompt Engineering
        prompt_template = PromptTemplate(
            input_variables=["query", "context"],
            template="""شما یک دستیار هوشمند هستید که به سوالات کاربران پاسخ می‌دهید.
سوال: {query}
اطلاعات مرتبط: {context}
لطفاً پاسخ دقیق و مختصر ارائه دهید."""
        )
        prompt = prompt_template.format(query=query, context=context)

        # استفاده از مدل gpt-4
        llm = ChatOpenAI(model_name="gpt-4", openai_api_key=os.getenv("My_huggingface_key"))
        response = llm.predict(prompt)

        final_response = f"پاسخ:\n{response}\n\nمنابع:\n"
        for doc in docs:
            final_response += f"- {doc.metadata.get('source', 'نامشخص')}, صفحه {doc.metadata.get('page', 'نامشخص')}\n"
        return final_response
    except Exception as e:
        logger.error(f"خطا در پاسخ‌گویی بر اساس سند: {e}")
        return f"❌ خطایی رخ داده است: {e}"

# رابط کاربری با Gradio
with gr.Blocks() as demo:
    gr.Markdown("# هوش مصنوعی همراه کارشناسان توزیع برق ایران")
    query = gr.Textbox(label="سوال خود را بپرسید", lines=2)
    response = gr.Textbox(label="پاسخ", lines=10, interactive=False)
    submit_btn = gr.Button("ارسال سوال")

    submit_btn.click(chat_with_doc, inputs=[query], outputs=[response])

demo.launch()