# rag_streamlit_app.py

import streamlit as st
import os
import warnings
import re
from dotenv import load_dotenv
from chat_logic import setup_default_rag, OPENAI_KEY # Import core logic

# Suppress LangChain and other warnings for a clean Streamlit app
warnings.filterwarnings("ignore")
load_dotenv()

# --- Configuration ---
st.set_page_config(page_title="Ring App RAG Chatbot", layout="wide")

# --- Initialize Session State ---
if 'chain' not in st.session_state:
    st.session_state.chain = None
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'memory' not in st.session_state:
    st.session_state.memory = None
if 'openai_api_key' not in st.session_state:
    st.session_state.openai_api_key = OPENAI_KEY


# --- Functions for UI Actions ---

def clear_chat_history():
    """Resets the chat history and the memory buffer."""
    if st.session_state.memory:
        st.session_state.memory.clear()
    st.session_state.chat_history = []
    st.toast("Chat history cleared!", icon="🧹")

def initialize_rag_system():
    """Initializes the RAG chain using the hardcoded default file."""
    if st.session_state.openai_api_key:
        with st.spinner("Setting up the Ring App knowledge base..."):
            try:
                model = "gpt-4-turbo" 
                
                # CALL THE NEW DEFAULT SETUP FUNCTION
                chain, memory = setup_default_rag(st.session_state.openai_api_key, model)
                
                st.session_state.chain = chain
                st.session_state.memory = memory
                st.session_state.chat_history = []
                st.toast("Ring App knowledge base loaded and chatbot ready!", icon="✅")
            except FileNotFoundError as e:
                 st.error(f"FATAL ERROR: {e}. Please ensure 'default_rag_file.pdf' is in the script directory.")
                 st.session_state.chain = None
            except Exception as e:
                st.error(f"Error setting up RAG system: {e}")
                st.session_state.chain = None
                st.session_state.memory = None
    elif not st.session_state.openai_api_key:
         st.error("Please enter your OpenAI API Key in the sidebar.")


def generate_response(prompt):
    """Invokes the RAG chain with the user's prompt."""
    if st.session_state.chain:
        try:
            # Invoke the chain
            response = st.session_state.chain.invoke({"question": prompt})
            answer = response.get("answer", "Sorry, I couldn't find an answer based only on the Ring App document.")
            
            # Clean response logic
            answer = re.sub(r'\\n|\r|\n', ' ', answer)
            answer = re.sub(r'(Sources?:\s*.+$)', '', answer, flags=re.IGNORECASE)
            answer = re.sub(r'\[[^\]]*\]|\([^\)]*\)', '', answer)
            answer = re.sub(r'[*_#>`~\-]{1,}', ' ', answer)
            answer = re.sub(r'\s{2,}', ' ', answer).strip()
            
            # Update chat history state
            st.session_state.chat_history.append({"role": "user", "content": prompt})
            st.session_state.chat_history.append({"role": "assistant", "content": answer})
            
            return answer

        except Exception as e:
            st.error(f"An error occurred during the conversation: {e}")
            return "Sorry, there was an error processing your request."
    else:
        return "Please initialize the chatbot using the button in the sidebar."


# --- Streamlit UI Layout ---

st.title("Ring App Support Chatbot")
st.markdown("This RAG system is pre-loaded with knowledge about the **Ring Doorbell App**")

# Sidebar for configuration
with st.sidebar:
    st.header("Configuration")
    
    # API Key Input
    st.session_state.openai_api_key = st.text_input(
        "OpenAI API Key", 
        value=st.session_state.openai_api_key, 
        type="password", 
        help="Required to use OpenAI embeddings and models."
    )
    
    st.markdown("---")
    
    # Initialization Button
    if st.button("Initialize Chatbot", type="primary"):
        initialize_rag_system()
        
    st.caption("The chatbot will only answer from the pre-loaded Ring App documentation.")
    
    st.markdown("---")
    
    # Reset Button
    if st.button("Clear History", help="Clears conversation memory and chat display."):
        clear_chat_history()
        
    # Check if the system is initialized and ready
    if st.session_state.chain:
        st.success("System Ready! Ask a question below.")


# --- Main Chat Interface ---

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

# Initial state prompt
if not st.session_state.chain and not st.session_state.chat_history:
    st.info("Click **Initialize Chatbot** in the sidebar to load the default Ring App knowledge base.")
    st.stop()


# Chat input box
if prompt := st.chat_input("Ask a question about Ring App setup, dashboard, or history..."):
    # Immediately display user message
    with st.chat_message("user"):
        st.write(prompt)
    
    # Generate and display assistant response
    with st.chat_message("assistant"):
        with st.spinner("Thinking..."):
            response_text = generate_response(prompt)
            st.write(response_text)