import streamlit as st

from rag.pdf_loader   import load_pdf
from rag.chunker      import chunk_text
from rag.embedder     import embed_chunks, embed_query
from rag.vector_store import build_index, search_index
from rag.llm          import ask_llm

# ── Page configuration ────────────────────────────────────────
# This must be the very first Streamlit command in the script
st.set_page_config(
    page_title = "RAG Assistant",
    page_icon  = "📄",
    layout     = "wide"
)

# ── Initialize session state ──────────────────────────────────
# session_state persists data between Streamlit reruns
# We set default values here if they do not exist yet
if "chunks"       not in st.session_state: st.session_state.chunks       = None
if "index"        not in st.session_state: st.session_state.index        = None
if "chat_history" not in st.session_state: st.session_state.chat_history = []
if "pdf_name"     not in st.session_state: st.session_state.pdf_name     = None

# ── Sidebar ───────────────────────────────────────────────────
with st.sidebar:
    st.title("📄 RAG Assistant")
    st.markdown("Chat with any PDF using AI — fully local embeddings, no data sent except your question.")
    st.divider()

    # Groq API key input
    # type="password" hides the key so it is not visible on screen
    st.subheader("Step 1 — Enter API Key")
    api_key = st.text_input(
        "Groq API Key",
        type        = "password",
        placeholder = "gsk_..."
    )
    st.caption("Get your free key at console.groq.com")

    st.divider()

    # PDF upload section
    st.subheader("Step 2 — Upload a PDF")
    uploaded_file = st.file_uploader(
        "Choose a PDF file",
        type = ["pdf"]
    )

    # Process the PDF when user uploads one
    if uploaded_file is not None:

        # Only re-process if it is a new PDF (not the same one again)
        if st.session_state.pdf_name != uploaded_file.name:

            with st.spinner("Reading and indexing your PDF..."):

                # Save the uploaded file temporarily to disk
                # Streamlit gives us the file as bytes, we write it to disk
                # so our load_pdf() function can open it with fitz
                temp_path = f"/tmp/{uploaded_file.name}"
                with open(temp_path, "wb") as f:
                    f.write(uploaded_file.read())

                # Run the full RAG pipeline on the PDF
                pages  = load_pdf(temp_path)                          # extract text
                chunks = chunk_text(pages, chunk_size=500, overlap=50) # split into chunks
                chunks = embed_chunks(chunks)                          # convert to vectors
                index  = build_index(chunks)                           # build FAISS index

                # Save everything to session state so it persists
                st.session_state.chunks   = chunks
                st.session_state.index    = index
                st.session_state.pdf_name = uploaded_file.name

                # Clear old chat history when a new PDF is loaded
                st.session_state.chat_history = []

            st.success(f"Ready! Loaded {len(pages)} pages, {len(chunks)} chunks.")

    # Show currently loaded PDF name
    if st.session_state.pdf_name:
        st.info(f"Current PDF: {st.session_state.pdf_name}")

    st.divider()
    st.caption("Built with PyMuPDF · SentenceTransformers · FAISS · Groq · Streamlit")

# ── Main area ─────────────────────────────────────────────────
st.title("Chat with your PDF")

# Show a welcome message if no PDF is loaded yet
if st.session_state.chunks is None:
    st.info("Upload a PDF in the sidebar to get started.")

else:
    # ── Chat history display ──────────────────────────────────
    # Loop through all previous messages and display them
    for message in st.session_state.chat_history:

        if message["role"] == "user":
            # Display user message on the right
            with st.chat_message("user"):
                st.write(message["content"])

        elif message["role"] == "assistant":
            # Display assistant answer on the left
            with st.chat_message("assistant"):
                st.write(message["content"])

                # Show the source pages used to generate this answer
                if "sources" in message:
                    with st.expander("View sources"):
                        for source in message["sources"]:
                            st.markdown(f"**Page {source['page_num']}**")
                            st.caption(source["text"][:300] + "...")
                            st.divider()

    # ── Question input ────────────────────────────────────────
    # st.chat_input stays fixed at the bottom of the page
    query = st.chat_input("Ask a question about your PDF...")

    if query:

        # Check that API key has been entered
        if not api_key:
            st.warning("Please enter your Groq API key in the sidebar.")

        else:
            # Show the user message immediately
            with st.chat_message("user"):
                st.write(query)

            # Save user message to chat history
            st.session_state.chat_history.append({
                "role":    "user",
                "content": query
            })

            # Generate the answer
            with st.chat_message("assistant"):
                with st.spinner("Thinking..."):

                    # Step 1 — convert question to vector
                    query_embedding = embed_query(query)

                    # Step 2 — find top 3 most relevant chunks
                    relevant_chunks = search_index(
                        st.session_state.index,
                        st.session_state.chunks,
                        query_embedding,
                        top_k = 3
                    )

                    # Step 3 — get answer from LLM
                    answer = ask_llm(query, relevant_chunks, api_key)

                # Display the answer
                st.write(answer)

                # Display the sources in an expandable section
                with st.expander("View sources"):
                    for source in relevant_chunks:
                        st.markdown(f"**Page {source['page_num']}**")
                        st.caption(source["text"][:300] + "...")
                        st.divider()

            # Save assistant answer + sources to chat history
            st.session_state.chat_history.append({
                "role":    "assistant",
                "content": answer,
                "sources": relevant_chunks
            })