import os
import streamlit as st
from dotenv import load_dotenv
import PyPDF2
import numpy as np
from openai import OpenAI
import openai  # لا نستورد error مباشرة

# Load environment
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

st.set_page_config(
    page_title="📚 المكتبة المريدية الذكية",
    page_icon="🚀",
    layout="wide"
)

# RTL + Styling
st.markdown("""
<style>
body { direction: rtl; font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; background-color: #f9fafb; }
.stButton>button { border-radius: 12px; background-color: #2563eb; color: white; font-weight: bold; }
.stTextInput>div>div>input { border-radius: 10px; }
h1, h2, h3 { color: #1e3a8a; }
.stChatMessage { direction: rtl; }
.stExpanderHeader { font-weight: bold; color: #1e40af; }
</style>
""", unsafe_allow_html=True)

# Sidebar
with st.sidebar:
    st.header("⚙️ الإعدادات / Settings")
    language = st.radio("اختر اللغة / Choose Language / Choisir la langue:", ["العربية", "English", "Français"])
    temperature = st.slider("Creativity / الإبداع", 0.0, 1.0, 0.2)
    st.markdown("---")
    st.info("قم برفع ملفات PDF الخاصة بك وابدأ بالبحث")

# PDF Upload
uploaded_files = st.file_uploader("📄 Upload PDF(s) / رفع ملفات PDF", type="pdf", accept_multiple_files=True)

# Simple VectorStore
class SimpleVectorStore:
    def __init__(self):
        self.embeddings = []
        self.documents = []

    def add_document(self, content, embedding):
        self.documents.append(content)
        self.embeddings.append(embedding)

    def query(self, query_embedding, top_k=4):
        if not self.embeddings:
            return []
        sims = np.dot(np.array(self.embeddings), query_embedding)
        top_indices = sims.argsort()[-top_k:][::-1]
        return [self.documents[i] for i in top_indices]

# Extract text
def extract_text_from_pdf(file):
    reader = PyPDF2.PdfReader(file)
    text = ""
    for page in reader.pages:
        text += page.extract_text() + "\n"
    return text

# Build VectorStore
@st.cache_resource
def build_vectorstore(files):
    vectorstore = SimpleVectorStore()
    client = OpenAI(api_key=OPENAI_API_KEY)
    for file in files:
        text = extract_text_from_pdf(file)
        chunks = [text[i:i+1000] for i in range(0, len(text), 1000)]
        for chunk in chunks:
            try:
                embedding = client.embeddings.create(
                    input=chunk,
                    model="text-embedding-3-small"
                )['data'][0]['embedding']
                vectorstore.add_document(chunk, np.array(embedding))
            except Exception as e:  # بديل عن RateLimitError
                st.warning(f"⚠️ بعض المستندات لم تُعالج: {str(e)}")
    return vectorstore

# Chat Memory
if "messages" not in st.session_state:
    st.session_state.messages = []

# Main Logic
if uploaded_files:
    vectorstore = build_vectorstore(uploaded_files)
    client = OpenAI(api_key=OPENAI_API_KEY)

    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.markdown(msg["content"])

    if prompt := st.chat_input("اكتب سؤالك هنا / Type your question:"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            with st.spinner("🤖 جاري معالجة السؤال..."):
                try:
                    query_embedding = client.embeddings.create(
                        input=prompt,
                        model="text-embedding-3-small"
                    )['data'][0]['embedding']

                    top_docs = vectorstore.query(np.array(query_embedding), top_k=4)
                    context = "\n\n".join(top_docs)

                    response = client.chat.completions.create(
                        model="gpt-4o-mini",
                        messages=[
                            {"role": "system", "content": "أنت مساعد بحثي."},
                            {"role": "user", "content": f"استعن بالنصوص التالية للإجابة:\n{context}\n\nسؤال: {prompt}"}
                        ],
                        temperature=temperature
                    )
                    answer = response.choices[0].message.content
                    st.markdown(answer)

                    with st.expander("📚 المصادر / Sources"):
                        for i, doc in enumerate(top_docs):
                            st.write(f"المصدر {i+1}: {doc[:500]}...")

                    st.session_state.messages.append({"role": "assistant", "content": answer})
                except Exception as e:
                    st.error(f"⚠️ حدث خطأ أثناء الاتصال بـ OpenAI: {str(e)}")
else:
    st.warning("📂 الرجاء رفع ملف PDF واحد على الأقل للبدء.")