import os
import shutil
import logging
import gradio as gr
from langchain_community.document_loaders import PyPDFLoader  # تغییر به PyPDFLoader
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.chat_models import ChatOpenAI

# تنظیمات لاگ‌گیری
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# تنظیمات مسیرها
UPLOAD_FOLDER = "uploaded_files"
DATABASE_FILE = "/home/user/app/vector_database"  # بدون پسوند .pkl
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# بررسی و بارگذاری دیتابیس برداری
def load_database():
    try:
        if os.path.exists(DATABASE_FILE):
            embeddings = OpenAIEmbeddings(openai_api_key=os.getenv("My_huggingface_key"))
            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

def save_database(vector_db):
    try:
        vector_db.save_local(DATABASE_FILE)
        logger.info(f"✅ دیتابیس ذخیره شد: {DATABASE_FILE}")
    except Exception as e:
        logger.error(f"❌ خطا در ذخیره دیتابیس: {e}")
        raise e

global_vector_db = load_database()

# پردازش و ذخیره‌سازی فایل‌های PDF
def process_and_store_pdfs(file_paths):
    texts = []
    for file_path in file_paths:
        try:
            logger.info(f"در حال پردازش فایل: {file_path}")
            if file_path.endswith(".pdf"):
                # استفاده از PyPDFLoader به جای PDFMinerLoader
                loader = PyPDFLoader(file_path)
                documents = loader.load()
                texts.extend(documents)
                logger.info(f"✅ تعداد صفحات پردازش شده: {len(documents)}")
            else:
                logger.warning(f"فرمت فایل {file_path} پشتیبانی نمی‌شود.")
        except Exception as e:
            logger.error(f"خطا در پردازش فایل {file_path}: {e}")

    try:
        logger.info("در حال ایجاد embeddings و ذخیره‌سازی در FAISS...")
        openai_api_key = os.getenv("My_huggingface_key")

        if not openai_api_key:
            logger.error("❌ کلید API مقداردهی نشده است!")
            raise ValueError("❌ کلید OpenAI API یافت نشد.")

        embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
        vector_db = FAISS.from_documents(texts, embeddings)
        logger.info(f"✅ ذخیره‌سازی در FAISS با موفقیت انجام شد. تعداد اسناد: {vector_db.index.ntotal}")
        return vector_db
    except Exception as e:
        logger.error(f"خطا در ایجاد embeddings یا ذخیره‌سازی در FAISS: {e}")
        return None

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

        if not query.strip():
            return "❌ لطفاً سوال خود را وارد کنید."

        # بازیابی اسناد مرتبط
        retriever = global_vector_db.as_retriever(search_kwargs={"k": 5})
        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 "هیچ اطلاعات مرتبطی یافت نشد."

        # استفاده از مدل gpt-3.5-turbo به جای gpt-4 (در صورت عدم دسترسی به GPT-4)
        llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=os.getenv("My_huggingface_key"))
        
        # ساختار درخواست ساده‌تر
        prompt = f"""سوال: {query}\n\nاطلاعات مرتبط:\n{context}\n\nلطفاً به سوال پاسخ دهید:"""
        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}"

# ذخیره فایل آپلود شده
def save_uploaded_file(file):
    try:
        file_name = os.path.basename(file)
        file_path = os.path.join(UPLOAD_FOLDER, file_name)
        shutil.copy(file, file_path)

        global global_vector_db
        global_vector_db = process_and_store_pdfs([file_path])
        if global_vector_db:
            save_database(global_vector_db)

        return f"✅ فایل ذخیره شد: {file_path}"
    except Exception as e:
        return f"❌ خطا در ذخیره فایل: {e}"

# دباگ موقت
def debug_temp():
    debug_results = []

    # بررسی مقدار کلید API
    openai_api_key = os.getenv("My_huggingface_key")
    if openai_api_key:
        debug_results.append(f"✅ کلید API مقداردهی شده: {openai_api_key[:5]}********")
    else:
        debug_results.append("❌ کلید API مقداردهی نشده است!")

    # بررسی نسخه langchain_community
    try:
        import langchain_community
        version = langchain_community.__version__
        debug_results.append(f"✅ نسخه langchain_community: {version}")
    except Exception as e:
        debug_results.append(f"❌ خطا در بررسی نسخه langchain_community: {e}")

    # بررسی فایل‌های PDF در مسیر آپلود
    try:
        pdf_files = [f for f in os.listdir(UPLOAD_FOLDER) if f.endswith(".pdf")]
        if pdf_files:
            file_paths = [os.path.join(UPLOAD_FOLDER, f) for f in pdf_files]
            debug_results.append(f"✅ فایل‌های PDF در مسیر `{UPLOAD_FOLDER}`:")
            
            # بررسی تعداد صفحات هر فایل
            for file_path in file_paths:
                try:
                    loader = PyPDFLoader(file_path)  # استفاده از PyPDFLoader
                    documents = loader.load()
                    debug_results.append(f"   - {os.path.basename(file_path)}: {len(documents)} صفحه")
                except Exception as e:
                    debug_results.append(f"❌ خطا در بررسی صفحات {file_path}: {str(e)}")
        else:
            debug_results.append(f"❌ هیچ فایل PDF در مسیر `{UPLOAD_FOLDER}` وجود ندارد.")
    except Exception as e:
        debug_results.append(f"❌ خطا در بررسی پوشه `{UPLOAD_FOLDER}`: {e}")

    # بررسی وضعیت دیتابیس
    try:
        if os.path.exists(DATABASE_FILE):
            debug_results.append(f"✅ دیتابیس در مسیر `{DATABASE_FILE}` وجود دارد.")
            
            # بررسی تعداد اسناد ذخیره شده
            if global_vector_db:
                debug_results.append(f"   - تعداد اسناد در دیتابیس: {global_vector_db.index.ntotal}")
            else:
                debug_results.append("❌ دیتابیس در حافظه بارگذاری نشده است")
        else:
            debug_results.append(f"❌ دیتابیس در مسیر `{DATABASE_FILE}` وجود ندارد.")
    except Exception as e:
        debug_results.append(f"❌ خطا در بررسی دیتابیس: {e}")

    return "\n".join(debug_results)

# رابط کاربری با 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("ارسال سوال")

    file = gr.File(label="📎 آپلود فایل", file_types=[".pdf"])
    upload_status = gr.Textbox(label="وضعیت آپلود", interactive=False)
    upload_btn = gr.Button("ارسال فایل")

    debug_btn = gr.Button("🔍 بررسی دباگ موقت")
    debug_output = gr.Textbox(label="نتایج دباگ", lines=10, interactive=False)

    upload_btn.click(save_uploaded_file, inputs=[file], outputs=[upload_status])
    submit_btn.click(chat_with_doc, inputs=[query], outputs=[response])
    debug_btn.click(debug_temp, outputs=debug_output)

demo.launch()