import streamlit as st
import os
from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings   # FIXED: new package
from langchain_chroma import Chroma
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

st.title("📚 MCA Module AI Assistant")
st.markdown("Ask anything about your Computer Vision modules!")

groq_api_key = os.getenv("GROQ_API_KEY")
llm = ChatGroq(model="llama-3.2-3b-preview", groq_api_key=groq_api_key)

@st.cache_resource
def prepare_vectorstore():
    loader = PyPDFDirectoryLoader("mca_modules/")
    docs = loader.load()
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
    splits = text_splitter.split_documents(docs)
    
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    
    return Chroma.from_documents(
        documents=splits,
        embedding=embeddings,
        persist_directory="/tmp/chroma_db"   # FIXED: writable path on HF Spaces
    )

vectorstore = prepare_vectorstore()

template = """You are an expert AI Assistant. Use the provided context to answer the question.
Rules:
1. If the user asks for a comparison or specific categories, use a Markdown TABLE.
2. Use bolding for key terms.
3. If the answer isn't in the context, say "I cannot find this in your modules."

Context:
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
rag_chain = (
    {"context": vectorstore.as_retriever(), "question": RunnablePassthrough()}
    | prompt | llm | StrOutputParser()
)

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

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt_input := st.chat_input("What is covered in Module 1?"):
    st.session_state.messages.append({"role": "user", "content": prompt_input})
    with st.chat_message("user"):
        st.markdown(prompt_input)

    with st.chat_message("assistant"):
        response = rag_chain.invoke(prompt_input)
        st.markdown(response)
        st.session_state.messages.append({"role": "assistant", "content": response})