import streamlit as st
import os
import requests
import re  # To help clean up leading whitespace
import sys

# ✅ Redirect Hugging Face and Streamlit caches to writable folders
os.environ["HF_HOME"] = "/tmp/huggingface"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp/huggingface"
os.environ["STREAMLIT_CACHE_DIR"] = "/tmp/streamlit"
os.environ["XDG_CACHE_HOME"] = "/tmp"

# ✅ (Optional) If using LangChain
os.environ["LANGCHAIN_CACHE"] = "/tmp/langchain"

# Prefer pysqlite3 (if installed) for SQLite; otherwise fall back to stdlib sqlite3.
# This mirrors the previous behavior but is robust when pysqlite3 isn't available.
try:
    # Try to import the pysqlite3 package which provides a replacement sqlite3 module
    pysqlite3 = __import__('pysqlite3')
    # If imported, make it available under the name 'sqlite3' to satisfy code expecting
    # a sqlite3 module with pysqlite3's behavior.
    if 'sqlite3' not in sys.modules:
        sys.modules['sqlite3'] = pysqlite3
except ModuleNotFoundError:
    # pysqlite3 isn't installed; the standard library's sqlite3 will be used.
    import sqlite3  # noqa: F401

# Langchain and HuggingFace
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
# Workaround: some versions of the `groq` package (used by `langchain_groq`) may pass
# a `proxies` keyword into the underlying Client.__init__, which can cause
# TypeError if the installed groq Client doesn't accept that kwarg. To make the
# app more tolerant across versions, monkeypatch the groq Client constructor to
# silently drop `proxies` if present.
try:
    import groq._base_client as _groq_base
    _groq_client_init = getattr(_groq_base.Client, "__init__", None)
    if _groq_client_init is not None:
        def _client_init_no_proxies(self, *args, **kwargs):
            kwargs.pop('proxies', None)
            return _groq_client_init(self, *args, **kwargs)
        _groq_base.Client.__init__ = _client_init_no_proxies
except Exception:
    # If groq isn't installed or the internal structure differs, let the import
    # of ChatGroq attempt to initialize and raise its own errors. We don't
    # want to crash here just for the monkeypatch.
    pass

from langchain_groq import ChatGroq
from langchain.chains import RetrievalQA

# Load the .env file (if using it)
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="openai/gpt-oss-20b", 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>
    .subtitle {text-align: center; font-size: 18px; font-weight: bold;}
    .details {font-size: 14px; color: #bbbbbb; padding-left: 20px; margin-bottom: -10px;}
    .team {text-align: center; font-size: 20px; font-weight: bold; color: #777;}
    </style>
    """,
    unsafe_allow_html=True
)

with st.sidebar:
    # App title
    st.title('CSPC Conversational Agent')
    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="details">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)