import streamlit as st
import time
import os
import pickle
import numpy as np
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_together import TogetherEmbeddings
from langchain.chat_models import ChatOpenAI
from sklearn.metrics.pairwise import cosine_similarity

# ------------------ بارگذاری چانک‌ها و امبدینگ‌ها ------------------

@st.cache_resource
def load_chunks_and_embeddings():
    embeddings_file = 'embeddings.pkl'

    if os.path.exists(embeddings_file):
        st.success("✅ امبدینگ‌ها از فایل کش بارگذاری شد.")
        with open(embeddings_file, 'rb') as f:
            data = pickle.load(f)
        return data['chunk_texts'], data['chunk_embeddings'], data['embeddings_model']

    else:
        with st.spinner('📄 در حال پردازش PDF و ساخت امبدینگ‌ها...'):
            loader = PyPDFLoader('test1.pdf')
            pages = loader.load()

            splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0)
            chunks = splitter.split_documents(pages)

            embeddings_model = TogetherEmbeddings(
                api_key="0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979"
            )

            chunk_texts = [chunk.page_content for chunk in chunks]

            # ساخت امبدینگ با progress bar
            progress = st.progress(0, text="در حال ساخت امبدینگ چانک‌ها...")
            chunk_embeddings = []
            for i, text in enumerate(chunk_texts):
                chunk_embeddings.append(embeddings_model.embed_query(text))
                progress.progress((i + 1) / len(chunk_texts))

            # ذخیره در فایل کش
            with open(embeddings_file, 'wb') as f:
                pickle.dump({
                    'chunk_texts': chunk_texts,
                    'chunk_embeddings': chunk_embeddings,
                    'embeddings_model': embeddings_model,
                }, f)

            st.success(f"✅ {len(chunk_texts)} چانک پردازش و ذخیره شد.")
            return chunk_texts, chunk_embeddings, embeddings_model

chunk_texts, chunk_embeddings, embeddings_model = load_chunks_and_embeddings()

# ------------------ ساخت مدل LLM ------------------

llm = ChatOpenAI(
    base_url="https://api.together.xyz/v1",
    api_key='0291f33aee03412a47fa5d8e562e515182dcc5d9aac5a7fb5eefdd1759005979',
    model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free"
)

# ------------------ پاسخ بر اساس ۱۰ چانک نزدیک ------------------

def answer_from_pdf(question):
    # ۱- ساخت امبدینگ سوال
    question_embedding = embeddings_model.embed_query(question)

    # ۲- محاسبه شباهت
    similarities = cosine_similarity(
        [question_embedding],
        chunk_embeddings
    )[0]

    # ۳- انتخاب ۱۰ چانک نزدیک
    top_indices = np.argsort(similarities)[-10:][::-1]
    selected_chunks = [chunk_texts[i] for i in top_indices]

    # ۴- ساخت پرامپت
    context = "\n\n".join(selected_chunks)
    prompt = f"""با توجه به متن زیر فقط به زبان فارسی پاسخ بده:

    متن:
    {context}

    سوال:
    {question}

    پاسخ:"""

    response = llm.invoke(prompt)
    return response.content

# ------------------ Chat Streamlit UI ------------------

st.title('📚 چت با PDF (با ۱۰ چانک نزدیک و کش شده)')

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 = answer_from_pdf(st.session_state.pending_prompt)
        answer = response.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