import streamlit as st
import os
import requests
from dotenv import load_dotenv  # Only needed if using a .env file
import re  # To help clean up leading whitespace


# 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.1)
    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"
)

# Custom CSS for styling
st.markdown(
    """
    <style>
    # .main {background-color: #f4f4f4;}
    .title {text-align: center; font-size: 30px; font-weight: bold; color: #ffffff;}
    .subtitle {text-align: center; font-size: 18px; font-weight: bold; color: #dddddd;}
    .category {font-size: 16px; font-weight: bold; color: #222;}
    .details {font-size: 14px; color: #bbbbbb; margin-left: 25px; margin-top: -10px; margin-bottom: 5px;}
    .details1 {font-size: 14px; color: #bbbbbb; margin-left: 25px; margin-top: -10px; margin-bottom: 10px;}
    .team {text-align: center; font-size: 14px; font-weight: bold; color: #777; margin-top: 20px;}
    .resources a {color: #0066cc; text-decoration: none; font-weight: bold;}
    .resources a:hover {color: #003366;}
    </style>
    """,
    unsafe_allow_html=True
)

with st.sidebar:
    # App title
    st.markdown('<p class="title">CSPC Conversational Agent</p>', unsafe_allow_html=True)
    st.markdown('<p class="subtitle">Your go-to assistant for the Citizen’s Charter of CSPC!</p>', unsafe_allow_html=True)
    
    # Categories
    st.markdown('''✔️**About CSPC:**''')
    st.markdown('<p class="details">History, Core Values, Mission and Vision</p>', unsafe_allow_html=True)
    
    st.markdown('''✔️**Admission & Graduation:**''')
    st.markdown('<p class="details">Apply, Requirements, Process, Graduation</p>', unsafe_allow_html=True)

    st.markdown('''✔️**Student Services:**''')
    st.markdown('<p class="details">Scholarships, Orgs, Facilities</p>', unsafe_allow_html=True)

    st.markdown('''✔️**Academics:**''')
    st.markdown('<p class="details">Degrees, Courses, Faculty</p>', unsafe_allow_html=True)

    st.markdown('''✔️**Officials:**''')
    st.markdown('<p class="details1">President, VPs, Deans, Admin</p>', unsafe_allow_html=True)

    # Links to resources
    st.markdown("### 🔗 Quick Access to Resources")
    st.markdown(
        """
        📄 [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/)  
        """,
        unsafe_allow_html=True
    )

# Store LLM generated responses
if "messages" not in st.session_state:
    st.session_state.chain = init_chain()
    st.session_state.messages = [{"role": "assistant", "content": "Hello! I am your Conversational Agent for the Citizens Charter of Camarines Sur Polytechnic Colleges (CSPC). How may I assist you today?"}]
    st.session_state.query_counter = 0  # Track the number of user queries
    st.session_state.conversation_history = ""  # Keep track of history for the LLM

def generate_response(prompt_input):
    try:
        # Retrieve vector database context using ONLY the current user input
        retriever = st.session_state.chain.retriever
        relevant_context = retriever.get_relevant_documents(prompt_input)  # Retrieve context only for the current prompt
        
        # Format the input for the chain with the retrieved context
        formatted_input = (
            f"You are a Conversational Agent for the Citizens Charter of Camarines Sur Polytechnic Colleges (CSPC). "
            f"Your purpose is to provide accurate and helpful information about CSPC's policies, procedures, and services as outlined in the Citizens Charter. "
            f"When responding to user queries:\n"
            f"1. Always prioritize information from the provided context (Citizens Charter or other CSPC resources).\n"
            f"2. Be concise, clear, and professional in your responses.\n"
            f"3. If the user's question is outside the scope of the Citizens Charter, politely inform them and suggest relevant resources or departments they can contact.\n\n"
            f"Context:\n"
            f"{' '.join([doc.page_content for doc in relevant_context])}\n\n"
            f"Conversation:\n{st.session_state.conversation_history}user: {prompt_input}\n"
        )
    
        # Invoke the RetrievalQA chain directly with the formatted input
        res = st.session_state.chain.invoke({"query": formatted_input})
        
        # Process the response text
        result_text = res['result']
        
        # Clean up prefixing phrases and capitalize the first letter
        if result_text.startswith('According to the provided context, '):
            result_text = result_text[35:].strip()
        elif result_text.startswith('Based on the provided context, '):
            result_text = result_text[31:].strip()
        elif result_text.startswith('According to the provided text, '):
            result_text = result_text[34:].strip()
        elif result_text.startswith('According to the context, '):
            result_text = result_text[26:].strip()
        
        # Ensure the first letter is uppercase
        result_text = result_text[0].upper() + result_text[1:] if result_text else result_text
    
        # Extract and format sources (if available)
        sources = []
        for doc in relevant_context:
            source_path = doc.metadata.get('source', '')
            formatted_source = source_path[122:-4] if source_path else "Unknown source"
            sources.append(formatted_source)
        
        # Remove duplicates and combine into a single string
        unique_sources = list(set(sources))
        source_list = ", ".join(unique_sources)
        
        # # Combine response text with sources
        # result_text += f"\n\n**Sources:** {source_list}" if source_list else "\n\n**Sources:** None"
        
        # Update conversation history
        st.session_state.conversation_history += f"user: {prompt_input}\nassistant: {result_text}\n"
    
        return result_text

    except Exception as e:
        # Handle rate limit or other errors gracefully
        if "rate_limit_exceeded" in str(e).lower():
            return "⚠️ Rate limit exceeded. Please clear the chat history and try again."
        else:
            return f"❌ An error occurred: {str(e)}"


# 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..."):
    # Increment query counter
    st.session_state.query_counter += 1
    # 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})

    # Check if query counter has reached the limit
    if st.session_state.query_counter >= 10:
        st.sidebar.warning("Conversation context has been reset after 10 queries.")
        st.session_state.query_counter = 0  # Reset the counter
        st.session_state.conversation_history = ""  # Clear conversation history for the LLM

# Clear chat history function
def clear_chat_history():
    # Clear chat messages (reset the assistant greeting)
    st.session_state.messages = [{"role": "assistant", "content": "Hello! I am your Conversational Agent for the Citizens Charter of Camarines Sur Polytechnic Colleges (CSPC). How may I assist you today?"}]
    
    # Reinitialize the chain to clear any stored history (ensures it forgets previous user inputs)
    st.session_state.chain = init_chain()

    # Clear the query counter and conversation history
    st.session_state.query_counter = 0
    st.session_state.conversation_history = ""

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

# Footer
st.sidebar.markdown('<p class="team">Developed by Team XceptionNet</p>', unsafe_allow_html=True)