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
from langchain.indexes import VectorstoreIndexCreator


@st.cache_resource
def load_chunks_and_embeddings():
    with st.spinner("در حال بارگذاری فایل و آماده‌سازی... لطفاً صبور باشید 🙏"):
        progress_bar = st.progress(0, text="در حال بارگذاری فایل PDF...")

        pdf_loader = PyPDFLoader('test1.pdf')
        pages = pdf_loader.load()
        progress_bar.progress(30, text="صفحات PDF بارگذاری شد. در حال ایجاد مدل برداری...")

        embeddings = TogetherEmbeddings(
            api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
        )
        progress_bar.progress(60, text="مدل Embedding ساخته شد. در حال ایجاد ایندکس...")

        index = VectorstoreIndexCreator(
            embedding=embeddings,
            text_splitter=RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0)
        ).from_loaders([pdf_loader])

        progress_bar.progress(100, text="بارگذاری کامل شد! ✅")
    return index
llm = ChatOpenAI(
    base_url="https://api.together.xyz/v1",
    api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979',
    model="meta-llama/Llama-3-70B-Instruct-Turbo-Free"
)

index = load_chunks_and_embeddings()

chain = RetrievalQA.from_chain_type(
  llm=llm,
  chain_type='stuff',
  retriever=index.vectorstore.as_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