# app.py
import streamlit as st
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains.question_answering import load_qa_chain
from langchain.llms import HuggingFacePipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import tempfile

# --- Streamlit page config ---
st.set_page_config(page_title="📚 Multi PDF Chatbot", layout="wide", page_icon="🤖")

# --- Header ---
st.markdown(
    "<h1 style='text-align: center; color: #2F4F4F;'>📚 Multi-PDF Chat Agent 🤖</h1>",
    unsafe_allow_html=True
)

# --- Sidebar Styling ---
st.markdown(
    """
    <style>
    section[data-testid="stSidebar"] {
        background-color: #111827;
        padding: 30px 20px;
    }
    .profile-img-container {
        display: flex;
        justify-content: center;
        margin-bottom: 15px;
    }
    .profile-img-container img {
        border-radius: 12px;
        height: 100px;
        width: 100px;
        object-fit: cover;
        border: 2px solid #4ade80;
        box-shadow: 0 0 10px rgba(0,0,0,0.5);
    }
    .upload-section {
        background-color: #1f2937;
        padding: 20px;
        border-radius: 12px;
        margin-top: 20px;
        margin-bottom: 30px;
    }
    .upload-section h4 {
        color: #facc15;
        margin-bottom: 15px;
    }
    .footer {
        text-align: center;
        font-size: 13px;
        color: #9ca3af;
    }
    .footer a {
        color: #facc15;
        text-decoration: none;
    }
    .footer a:hover {
        color: #ffffff;
    }
    </style>
    """,
    unsafe_allow_html=True
)

# --- Sidebar Layout ---
with st.sidebar:
    # Profile image
    col1, col2, col3 = st.columns([1, 2, 1])
    with col2:
        st.image("assets/img/main.png", width=100)

    # Upload PDFs
    st.markdown('<div class="upload-section">', unsafe_allow_html=True)
    st.markdown("#### 📁 Upload PDF Files")
    pdf_docs = st.file_uploader("Drag and drop your PDFs here", accept_multiple_files=True, label_visibility="collapsed")
    st.markdown('</div>', unsafe_allow_html=True)

    # Footer
    st.markdown(
        """
        <div class="footer">
            Built with ❤️ by <a href="https://github.com/Danish7861" target="_blank">Danish Shahzad</a>
        </div>
        """,
        unsafe_allow_html=True
    )

# --- Main Logic ---

# Store FAISS vector store in session state
if "vectorstore" not in st.session_state:
    st.session_state.vectorstore = None

if pdf_docs:
    if st.button("📤 Submit & Process"):
        with st.spinner("Processing PDFs..."):
            all_text = ""
            for pdf_file in pdf_docs:
                pdf_reader = PdfReader(pdf_file)
                for page in pdf_reader.pages:
                    all_text += page.extract_text() + "\n"

            # Split text into chunks
            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
            chunks = splitter.split_text(all_text)

            # Create embeddings
            embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
            st.session_state.vectorstore = FAISS.from_texts(chunks, embeddings)
            st.success("✅ PDF content indexed successfully!")

# --- Load local LLM ---
@st.cache_resource(show_spinner=False)
def load_local_llm():
    tokenizer = AutoTokenizer.from_pretrained("TheBloke/guanaco-7B-GGML")
    model = AutoModelForCausalLM.from_pretrained("TheBloke/guanaco-7B-GGML")
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=512)
    return HuggingFacePipeline(pipeline=pipe)

llm = load_local_llm()

# --- Question input ---
user_question = st.text_input("🔍 Ask something from your uploaded PDFs:")
if user_question:
    if st.session_state.vectorstore is None:
        st.warning("⚠️ Please upload and process PDFs first.")
    else:
        with st.spinner("Thinking... 💭"):
            relevant_docs = st.session_state.vectorstore.similarity_search(user_question)
            chain = load_qa_chain(llm, chain_type="stuff")
            answer = chain.run(input_documents=relevant_docs, question=user_question)
            st.success("Answer:")
            st.write(answer)

# --- Fixed footer ---
st.markdown(
    """
    <div style="position: fixed; bottom: 0; width: 100%; background-color: #2F4F4F; padding: 10px; color: white; text-align: center; font-size: 13px;">
        📄 Multi-PDF Chatbot | Powered by LangChain, FAISS & Local AI
    </div>
    """,
    unsafe_allow_html=True
)