import streamlit as st
import tempfile
from langchain_huggingface import ChatHuggingFace, HuggingFaceEmbeddings, HuggingFaceEndpoint
from langchain_core.prompts import PromptTemplate
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.document_loaders import PyPDFLoader
from langchain.vectorstores import FAISS
from langchain_core.output_parsers import StrOutputParser
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage
from dotenv import load_dotenv

load_dotenv()

# -----------------------
# Initialize models
# -----------------------
llm = HuggingFaceEndpoint(
    model='meta-llama/Llama-3.1-8B-Instruct',
    task='text-generation'
)
chat_model = ChatHuggingFace(llm=llm)
embed_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-mpnet-base-v2')

# -----------------------
# Streamlit UI
# -----------------------
st.set_page_config(page_title="Talk with your Documents", layout="wide")
st.sidebar.title("Settings")

chunk_size = st.sidebar.slider("Chunk Size", min_value=100, max_value=1000, value=300, step=50)
chunk_overlap = st.sidebar.slider("Chunk Overlap", min_value=0, max_value=200, value=25, step=5)

uploaded_files = st.sidebar.file_uploader(
    "Upload PDF(s)",
    type=["pdf"],
    accept_multiple_files=True
)

# -----------------------
# Initialize FAISS store in session
# -----------------------
if "vector_store" not in st.session_state:
    if uploaded_files:
        all_docs = []
        for uploaded_file in uploaded_files:
            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
                tmp_file.write(uploaded_file.read())
                tmp_file_path = tmp_file.name
            loader = PyPDFLoader(tmp_file_path)
            docs = loader.load()
            all_docs.extend(docs)

        splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        splitted_docs = splitter.split_documents(all_docs)
        st.session_state.vector_store = FAISS.from_documents(splitted_docs, embedding=embed_model)

# -----------------------
# Chat logic
# -----------------------
st.title("📄 Talk with your Documents - RAG Based")

if "messages" not in st.session_state:
    st.session_state.messages = [SystemMessage(content="You are a helpful assistant")]

query = st.text_input("Ask a question about your uploaded documents:")

if st.button("Ask"):
    if not uploaded_files:
        st.error("❌ Please upload at least one PDF file before asking a question!")
    elif not query.strip():
        st.warning("⚠️ Please enter a question to ask.")
    else:
        st.session_state.messages.append(HumanMessage(content=query))
        
        with st.spinner("Finding the answer..."):
            # Use vector store from session
            retriever = st.session_state.vector_store.as_retriever(search_kwargs={'k': 3})
            context_chunks = retriever.invoke(query)
            final_context = " ".join([chunk.page_content for chunk in context_chunks])
            
            # Build chat history string
            history_str = "\n".join([f"{type(m).__name__}: {m.content}" for m in st.session_state.messages])
            
            # Prepare prompt
            prompt = PromptTemplate(
                template="""
You are a helpful assistant.
Answer the user's question ONLY based on the context provided.
If the answer is not in the context, reply: 'This is not included in your File.'
Context: {context}
Question: {query}
Previous chat history:
{history}
"""
            )
            final_prompt = prompt.format(context=final_context, query=query, history=history_str)
            
            # Get answer
            parser = StrOutputParser()
            result = chat_model.invoke(final_prompt)
            answer = result.content
            st.session_state.messages.append(AIMessage(content=answer))

# -----------------------
# Display recent chat
# -----------------------
st.subheader("💬 Recent Chat")
for msg in st.session_state.messages[1:]:
    if isinstance(msg, HumanMessage):
        st.markdown(f"**You:** {msg.content}")
    elif isinstance(msg, AIMessage):
        st.markdown(f"**AI:** {msg.content}")