import streamlit as st

from src.chunking import chunk_documents
from src.loader import load_all_pdfs
from src.memory import add_to_memory, get_memory
from src.rag_pipelines import build_context, generate_answer, retrieve_documents
from src.vector_store import create_vector_store

st.set_page_config(page_title="AI Pharma Support", layout="wide")
st.title("AI Pharmaceutical Support Assistant")


@st.cache_resource
def load_system():
    docs = load_all_pdfs()
    if not docs:
        raise ValueError("No PDF documents found in the data/ directory.")

    chunks = chunk_documents(docs)
    if not chunks:
        raise ValueError("PDFs were loaded, but no text chunks were created.")

    return create_vector_store(chunks)


if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

try:
    db = load_system()
except Exception as exc:
    st.error(
        "App initialization failed. Check OPENAI_API_KEY, dependencies, and data/*.pdf files."
    )
    st.exception(exc)
    st.stop()

user_input = st.text_input("Ask your pharma support question:")

if user_input:
    docs = retrieve_documents(user_input, db)
    context, sources = build_context(docs)
    memory = get_memory()
    answer = generate_answer(user_input, context, sources, memory)

    add_to_memory(user_input, answer)
    st.session_state.chat_history.append(("You", user_input))
    st.session_state.chat_history.append(("AI", answer))

for role, message in st.session_state.chat_history:
    if role == "You":
        st.markdown(f"**You:** {message}")
    else:
        st.markdown(f"**AI:** {message}")