import streamlit as st
import time
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_together import TogetherEmbeddings
from langchain.vectorstores import FAISS
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA

# --- 📄 ساخت امبدینگ‌ها با batch 50 تایی
def batch_embed(texts, embeddings_model, batch_size=50):
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        embs = embeddings_model.embed_documents([doc.page_content for doc in batch])
        all_embeddings.extend(embs)
    return all_embeddings

@st.cache_resource
def load_chunks_and_embeddings():
    pdf_loader = PyPDFLoader('test1.pdf')
    pages = pdf_loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0)
    docs = text_splitter.split_documents(pages)

    embeddings = TogetherEmbeddings(
        api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
    )
    
    vectorstore = None  # هنوز نساختیم

    # پروگرس بار
    progress = st.progress(0, text="🔄 در حال پردازش چانک‌ها...")
    total = len(docs)

    batch_size = 50
    for i in range(0, total, batch_size):
        batch_docs = docs[i:i+batch_size]
        embeddings_batch = embeddings.embed_documents([doc.page_content for doc in batch_docs])

        if vectorstore is None:
            vectorstore = FAISS.from_embeddings(embeddings_batch, batch_docs)
        else:
            vectorstore.add_embeddings(embeddings_batch, batch_docs)

        progress.progress(min((i+batch_size)/total, 1.0))

    progress.empty()
    return vectorstore

# --- 🛠️ آماده کردن دیتابیس
with st.spinner("📚 در حال بارگذاری فایل و ساخت امبدینگ‌ها... لطفا صبور باشید"):
    vectorstore = load_chunks_and_embeddings()

# --- 🤖 آماده سازی مدل LLM
llm = ChatOpenAI(
    base_url="https://api.together.xyz/v1",
    api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979',
    model="meta-llama/Llama-3-70B-Instruct-Turbo-Free"
)

retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 10})

chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type='stuff',
    retriever=retriever,
    input_key='question'
)

# --- 💬 چت بات
if 'messages' not in st.session_state:
    st.session_state.messages = []

if 'pending_prompt' not in st.session_state:
    st.session_state.pending_prompt = None

st.title("📄🤖 دستیار PDF شما")

# نمایش تاریخچه گفتگو
for msg in st.session_state.messages:
    with st.chat_message(msg['role']):
        st.markdown(f"🗨️ {msg['content']}", unsafe_allow_html=True)

prompt = st.chat_input("سوالی از PDF داری؟")

if prompt:
    st.session_state.messages.append({'role': 'user', 'content': prompt})
    st.session_state.pending_prompt = prompt
    st.rerun()

if st.session_state.pending_prompt:
    with st.chat_message('ai'):
        thinking = st.empty()
        thinking.markdown("🤖 در حال فکر کردن...")

        # اجرای جستجو در ایندکس
        response = chain.run(f'فقط به زبان فارسی جواب بده. سوال: {st.session_state.pending_prompt}')
        answer = response.split("Helpful Answer:")[-1].strip()
        if not answer:
            answer = "متأسفم، اطلاعات دقیقی در این مورد ندارم."

        thinking.empty()
        full_response = ""
        placeholder = st.empty()

        for word in answer.split():
            full_response += word + " "
            placeholder.markdown(full_response + "▌")
            time.sleep(0.03)

        placeholder.markdown(full_response)
        st.session_state.messages.append({'role': 'ai', 'content': full_response})
        st.session_state.pending_prompt = None