index.py

import os
import streamlit as st
from dotenv import load_dotenv

from app.loaders import load_and_chunk_files
from app.vectorstore import store_chunks, get_vectorstore, get_bm25_retriever
from app.chain import build_llm_chain, retrieve_hybrid_docs, rerank_documents
from app.pdf_handler import upload_pdfs, save_uploaded_files, list_uploaded_files, delete_uploaded_file

# Load environment variables
load_dotenv()

st.set_page_config(page_title="📄 Chat with PDF", layout="wide")
st.title("AI PDF RAG Agent: Chat with your PDFs")

UPLOAD_DIR = "uploaded_files"
os.makedirs(UPLOAD_DIR, exist_ok=True)

# STEP 1: Sidebar upload + config
uploaded_files, submitted, vectorstore_backend, chunk_size, overlap = upload_pdfs()


vectorstore = None
all_chunks = []

# STEP 2: Preview chunks for uploaded PDFs (optional)
if uploaded_files and not submitted:
    st.sidebar.subheader("📑 Chunk Preview (First 5 Chunks of Each PDF)")
    import tempfile
    preview_file_paths = []
    for file in uploaded_files:
        tmp_path = os.path.join(tempfile.gettempdir(), file.name)
        with open(tmp_path, "wb") as f:
            f.write(file.read())
        preview_file_paths.append(tmp_path)

    preview_chunks = load_and_chunk_files(
        preview_file_paths, chunk_size=chunk_size, chunk_overlap=overlap
    )
    for i, doc in enumerate(preview_chunks[:5]):
        st.sidebar.markdown(f"**Chunk {i+1}**: {doc.page_content[:200]}...")

# STEP 3: Load + Index PDFs
if uploaded_files and submitted:
    file_paths = save_uploaded_files(uploaded_files)
    st.sidebar.success(f"Uploaded {len(file_paths)} file(s)")

    with st.spinner("🔍 Indexing PDFs with FAISS..."):
        all_chunks = load_and_chunk_files(
            file_paths,
            chunk_size=chunk_size,
            chunk_overlap=overlap
        )
        st.success(f"✅ Indexed {len(all_chunks)} chunks from uploaded PDFs!")

        # Store chunks in FAISS
        vectorstore = store_chunks(all_chunks)
        bm25 = get_bm25_retriever(all_chunks)

# STEP 3b: Clear vectorstore
if st.sidebar.button("🧹 Clear Vectorstore"):
    if vectorstore:
        vectorstore.delete_all()
        all_chunks = []
        st.sidebar.success("Vectorstore cleared!")
        st.experimental_rerun()

# STEP 3c: Manage uploaded PDFs
st.sidebar.subheader("📄 Manage Uploaded PDFs")
existing_files = list_uploaded_files()
for f in existing_files:
    col1, col2 = st.sidebar.columns([4, 1])
    col1.markdown(f"- {f}")
    if col2.button("🗑️ Delete", key=f):
        delete_uploaded_file(f)
        st.experimental_rerun()

# STEP 4: Question input
if vectorstore:
    st.header("💬 Ask a question")
    query = st.text_input("Enter your question", key="user_question")

    if query:
        retrieved_docs = retrieve_hybrid_docs(query, vectorstore)
        reranked_docs = rerank_documents(query, retrieved_docs)
        chain = build_llm_chain()

        st.subheader("🤖 Answer:")
        response_container = st.empty()
        response = chain.invoke({"question": query, "docs": reranked_docs})
        response_container.markdown(response)

        st.sidebar.subheader("🔍 Retrieved Chunks")
        if reranked_docs:
            for i, doc in enumerate(reranked_docs):
                st.sidebar.markdown(f"**Chunk {i+1}**")
                st.sidebar.caption(doc.page_content[:400])
        else:
            st.sidebar.info("No chunks retrieved yet.")