app/main.py

"""
main.py
-------
Streamlit UI for the RAG chatbot.

Step 3 Enhancements:
- Score Threshold slider (0.0–1.0) in sidebar
- Streaming responses via st.write_stream() — no more spinner waiting
- Reset Knowledge Base button — deletes FAISS index, clears cache, resets state

All previous features retained:
- Upload PDF / TXT / DOCX / MD
- Top-K slider
- Clear Chat History button
- Conversation memory (last 6 messages)
- Source citations with similarity scores
- Cross-platform temp file handling
"""

import logging
import shutil
import sys
import tempfile
from pathlib import Path

import streamlit as st

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from app.chatbot import Chatbot
from app.config import APP_DESCRIPTION, APP_TITLE, VECTOR_DB_PATH
from components.document_loader import load_document
from components.embedder import HuggingFaceEmbedder
from components.text_splitter import split_documents
from components.vector_store import VectorStore

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ── Page config ───────────────────────────────────────────────────────────────
st.set_page_config(
    page_title=APP_TITLE,
    page_icon="🤖",
    layout="wide",
)

# ── Session state ─────────────────────────────────────────────────────────────
if "messages" not in st.session_state:
    st.session_state.messages = []

if "store_ready" not in st.session_state:
    st.session_state.store_ready = False

if "top_k" not in st.session_state:
    st.session_state.top_k = 4

if "score_threshold" not in st.session_state:
    st.session_state.score_threshold = 0.0


# ── Helper: initialise chatbot (cached) ───────────────────────────────────────
@st.cache_resource(show_spinner="Loading models …")
def get_chatbot() -> tuple:
    embedder = HuggingFaceEmbedder()
    store    = VectorStore(embedder=embedder, index_path=VECTOR_DB_PATH)
    loaded   = store.load()
    chatbot  = Chatbot(vector_store=store)
    return chatbot, loaded


# ── Sidebar ───────────────────────────────────────────────────────────────────
with st.sidebar:
    st.header("📂 Knowledge Base")
    st.caption("Upload documents and ingest them into the vector store.")

    uploaded_files = st.file_uploader(
        "Upload PDF / TXT / DOCX / MD",
        type=["pdf", "txt", "docx", "md"],
        accept_multiple_files=True,
    )

    if st.button("⚙️ Ingest Documents", use_container_width=True):
        if not uploaded_files:
            st.warning("Please upload at least one document first.")
        else:
            with st.spinner("Ingesting documents …"):
                chatbot_obj, _ = get_chatbot()
                all_chunks = []

                for uf in uploaded_files:
                    # Cross-platform temp dir (works on Windows + Linux/Mac)
                    tmp_path = Path(tempfile.gettempdir()) / uf.name
                    with open(tmp_path, "wb") as f:
                        f.write(uf.getbuffer())

                    try:
                        docs = load_document(tmp_path)
                        for d in docs:
                            d.metadata["source"] = uf.name
                        chunks = split_documents(docs)
                        all_chunks.extend(chunks)
                        st.success(f"✅ {uf.name} — {len(chunks)} chunks")
                    except Exception as exc:
                        st.error(f"❌ {uf.name}: {exc}")

                if all_chunks:
                    chatbot_obj.vector_store.build(all_chunks)
                    st.session_state.store_ready = True
                    st.success(f"Vector store built with {len(all_chunks)} total chunks.")

    st.divider()
    st.caption("Or pre-load documents by placing files in `data/raw/` and running `scripts/ingest.py`.")

    # ── Status badge ─────────────────────────────────────────────────────────
    chatbot_obj, preloaded = get_chatbot()
    ready = preloaded or st.session_state.store_ready

    if ready:
        st.success("📚 Knowledge base ready")
    else:
        st.warning("⚠️ No knowledge base loaded")

    # ── Top-K slider ──────────────────────────────────────────────────────────
    st.divider()
    top_k = st.slider(
        "Retrieved chunks (Top-K)",
        min_value=1,
        max_value=10,
        value=st.session_state.top_k,
        help="How many document chunks the retriever fetches per query. Higher = more context but slower.",
    )
    st.session_state.top_k = top_k

    # ── Score Threshold slider (Part A) ───────────────────────────────────────
    score_threshold = st.slider(
        "Score Threshold",
        min_value=0.0,
        max_value=1.0,
        value=st.session_state.score_threshold,
        step=0.05,
        help="Minimum similarity score a chunk must have to be used as context. "
             "0.0 = include everything. 0.5 = only confident matches.",
    )
    st.session_state.score_threshold = score_threshold

    # ── Clear Chat button ─────────────────────────────────────────────────────
    st.divider()
    if st.button("🗑️ Clear Chat History", use_container_width=True):
        st.session_state.messages = []
        st.rerun()

    # ── Reset Knowledge Base button (Part C) ──────────────────────────────────
    st.divider()
    if st.button("🔄 Reset Knowledge Base", use_container_width=True):
        index_path = Path(VECTOR_DB_PATH)
        if index_path.exists():
            shutil.rmtree(index_path)   # delete FAISS index folder + contents
        st.session_state.store_ready = False
        st.cache_resource.clear()       # force get_chatbot() to run fresh
        st.success("Knowledge base cleared. Upload new documents to start fresh.")
        st.rerun()


# ── Main chat UI ──────────────────────────────────────────────────────────────
st.title(f"🤖 {APP_TITLE}")
st.caption(APP_DESCRIPTION)
st.divider()

# Render chat history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])
        if msg.get("sources"):
            with st.expander("📄 Sources"):
                for src in msg["sources"]:
                    st.markdown(f"- `{src}`")

# Chat input
if prompt := st.chat_input("Ask a question about your documents …"):
    # Display user message
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # ── Streaming response (Part B) ───────────────────────────────────────────
    with st.chat_message("assistant"):
        top_k           = st.session_state.get("top_k", 4)
        score_threshold = st.session_state.get("score_threshold", 0.0)
        history         = st.session_state.messages[-6:] if st.session_state.messages else []

        chatbot_obj, _ = get_chatbot()

        token_stream, sources, _ = chatbot_obj.chat_stream(
            prompt,
            top_k=top_k,
            history=history,
            score_threshold=score_threshold,
        )

        if token_stream is None:
            # Early exit — store not ready or no relevant chunks found
            if not chatbot_obj.vector_store.is_ready:
                answer = "No documents have been ingested yet. Please upload documents first."
            else:
                answer = "I couldn't find any relevant information to answer your question."
            st.markdown(answer)
        else:
            # Stream tokens into UI — st.write_stream returns full text when done
            answer = st.write_stream(token_stream)

        if sources:
            with st.expander("📄 Sources"):
                for src in sources:
                    st.markdown(f"- `{src}`")

    # Save full answer to session state for chat history
    st.session_state.messages.append({
        "role":    "assistant",
        "content": answer,
        "sources": sources,
    })