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

# --------------------------------------------
# 1. بارگذاری پی دی اف و ساخت امبدینگ چانک‌ها (فقط یکبار و کش شده)
# --------------------------------------------
@st.cache_resource
def get_chunks_and_embeddings():
    pdf_loader = PyPDFLoader('test1.pdf')
    pages = pdf_loader.load_and_split(RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0))
    chunk_texts = [page.page_content for page in pages]

    embeddings_model = TogetherEmbeddings(
        api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
    )
    return chunk_texts, embeddings_model

# کش شده
chunk_texts, embeddings_model = get_chunks_and_embeddings()

# --------------------------------------------
# 2. ساختن امبدینگ چانک‌ها با پروگرس بار
# --------------------------------------------
st.title("📄 Chat with your PDF (با پی‌دی‌اف خودت حرف بزن!)")
st.subheader("در حال آماده‌سازی امبدینگ چانک‌ها...")

progress_bar = st.progress(0)
all_embeddings = []
batch_size = 128

for i in range(0, len(chunk_texts), batch_size):
    batch = chunk_texts[i:i+batch_size]
    embedded = embeddings_model.embed_documents(batch)
    all_embeddings.extend(embedded)

    progress_bar.progress(min((i + batch_size) / len(chunk_texts), 1.0))

st.success("✅ همه چانک‌ها آماده شدند!")

# --------------------------------------------
# 3. ساختن ایندکس FAISS از امبدینگ‌ها
# --------------------------------------------
vectorstore = FAISS.from_embeddings(all_embeddings, chunk_texts)

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

chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type='stuff',
    retriever=vectorstore.as_retriever(search_kwargs={"k": 10}),
    input_key='question'
)

# --------------------------------------------
# 5. چت بات Streamlit
# --------------------------------------------
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

# نمایش پیام‌های قبلی
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("چطور میتونم کمکت کنم؟")

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