import streamlit as st
import os
from config import OPENROUTER_API_KEY, HF_REPO_ID, HF_FILENAME
from rag import initialize_components, chat, reset_vector_store

def main():
    st.set_page_config(page_title="RAG Chatbot with ChromaDB", page_icon="🤖")
    st.title("🤖 RAG Chatbot with ChromaDB")
    st.markdown("Chat with traffic rules document from Hugging Face Hub!")
    
    # Check for OpenRouter API key
    if not OPENROUTER_API_KEY:
        st.error("Please set OPENROUTER_API_KEY environment variable")
        st.stop()
    
    # Initialize chat history
    if 'messages' not in st.session_state:
        st.session_state.messages = []
    
    # Sidebar for document management
    with st.sidebar:
        st.header("📄 Document Source")
        
        # Display document info
        st.info(f"""
        **Current Document:**
        - Repository: {HF_REPO_ID}
        - File: {HF_FILENAME}
        """)
        
        # Add reset button
        if st.button("Reset Vector Store"):
            if reset_vector_store():
                st.success("Vector store reset successfully!")
                st.experimental_rerun()
            else:
                st.info("Vector store does not exist.")
        
        # Add information about first-time model download
        st.info("""
        **First-time setup notice:**
        - The embedding model will be downloaded on first run (about 90MB).
        - If you see a message about 'Xet Storage' in the terminal, you can improve download speed by running:
          ```
          pip install hf_xet
          ```
        """)
    
    # Initialize components
    initialize_components()
    
    # Chat interface
    st.header("💬 Chat")
    
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])
    
    # Chat input
    if prompt := st.chat_input("Ask me anything about traffic rules..."):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.write(prompt)
        
        # Generate bot response
        with st.chat_message("assistant"):
            response = chat(prompt)
            st.write(response)
            st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    main()