import streamlit as st
import os
import requests
from dotenv import load_dotenv  # Only needed if using a .env file

# Langchain and HuggingFace
from langchain.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_groq import ChatGroq
from langchain.chains import RetrievalQA

# Load the .env file (if using it)
load_dotenv()
groq_api_key = os.getenv("GROQ_API_KEY")

# Load embeddings, model, and vector store
@st.cache_resource  # Singleton, prevent multiple initializations
def init_chain():
    
    model_kwargs = {'trust_remote_code': True}
    embedding = HuggingFaceEmbeddings(model_name='nomic-ai/nomic-embed-text-v1.5', model_kwargs=model_kwargs)
    llm = ChatGroq(groq_api_key=groq_api_key, model_name="llama3-70b-8192", temperature=0.2)
    vectordb = Chroma(persist_directory='updated_CSPCDB2', embedding_function=embedding)

    # Create chain
    chain = RetrievalQA.from_chain_type(llm=llm,
                                        chain_type="stuff",
                                        retriever=vectordb.as_retriever(k=5),
                                        return_source_documents=True)
    return chain

# Streamlit app layout
st.set_page_config(
    page_title="CSPC Citizens Charter Conversational Agent",
    page_icon="cspclogo.png"
)

with st.sidebar:
    st.title('CSPCean Conversational Agent')
    st.subheader('Ask anything CSPC Related here!')

    st.markdown('''**About CSPC:**
        History, Core Values, Mission and Vision''')

    st.markdown('''**Admission & Graduation:**
        Apply, Requirements, Process, Graduation''')

    st.markdown('''**Student Services:**
         Scholarships, Orgs, Facilities''')

    st.markdown('''**Academics:**
         Degrees, Courses, Faculty''')

    st.markdown('''**Officials:**
         President, VPs, Deans, Admin''')

    st.markdown('''
        Access the resources here:

        - [CSPC Citizen’s Charter](https://cspc.edu.ph/governance/citizens-charter/)
        - [About CSPC](https://cspc.edu.ph/about/)
        - [College Officials](https://cspc.edu.ph/college-officials/)
    ''')
    st.markdown('Team XceptionNet')

# Store LLM generated responses
if "messages" not in st.session_state:
    st.session_state.chain = init_chain()
    st.session_state.messages = [{"role": "assistant", "content": "How may I help you today?"}]

# Function for generating response using the last three conversations
def generate_response(prompt_input):
    # Initialize result
    result = ''
    
    # Prepare conversation history: get the last 3 user and assistant messages
    conversation_history = ""
    recent_messages = st.session_state.messages[-3:]  # Last 3 user and assistant exchanges (each exchange is 2 messages)

    for message in recent_messages:
        conversation_history += f"{message['role']}: {message['content']}\n"
    
    # Append the current user prompt to the conversation history
    conversation_history += f"user: {prompt_input}\n"

    # Invoke chain with the truncated conversation history
    res = st.session_state.chain.invoke(conversation_history)
    
    # Process response (as in the original code)
    if res['result'].startswith('According to the provided context, '):
        res['result'] = res['result'][35:]
        res['result'] = res['result'][0].upper() + res['result'][1:]
    elif res['result'].startswith('Based on the provided context, '):
        res['result'] = res['result'][31:]
        res['result'] = res['result'][0].upper() + res['result'][1:]
    elif res['result'].startswith('According to the provided text, '):
        res['result'] = res['result'][34:]
        res['result'] = res['result'][0].upper() + res['result'][1:]
    elif res['result'].startswith('According to the context, '):
        res['result'] = res['result'][26:]
        res['result'] = res['result'][0].upper() + res['result'][1:]
        
    
    # result += res['result']
    
    # # Process sources
    # result += '\n\nSources: '
    # sources = []
    # for source in res["source_documents"]:
    #     sources.append(source.metadata['source'][122:-4])  # Adjust as per your source format

    # sources = list(set(sources))  # Remove duplicates
    # source_list = ", ".join(sources)
    
    # result += source_list

    # return result, res['result'], source_list
    # return result, res['result']
    return res['result']

# Display chat messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])

# User-provided prompt for input box
if prompt := st.chat_input(placeholder="Ask a question..."):
    # Append user query to session state
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.write(prompt)

    # Generate and display placeholder for assistant response
    with st.chat_message("assistant"):
        message_placeholder = st.empty()  # Placeholder for response while it's being generated
        with st.spinner("Generating response..."):
            # Use conversation history when generating response
            response = generate_response(prompt)
        message_placeholder.markdown(response)  # Replace placeholder with actual response
        st.session_state.messages.append({"role": "assistant", "content": response})

# Clear chat history function
def clear_chat_history():
    # Clear chat messages (reset the assistant greeting)
    st.session_state.messages = [{"role": "assistant", "content": "How may I help you today?"}]
    
    # Reinitialize the chain to clear any stored history (ensures it forgets previous user inputs)
    st.session_state.chain = init_chain()

    # Clear any additional session state that might be remembering user inquiries
    if "recent_user_messages" in st.session_state:
        del st.session_state["recent_user_messages"]  # Clear remembered user inputs

st.sidebar.button('Clear Chat History', on_click=clear_chat_history)