import streamlit as st
import os
from chatbot import LearningChatbot
import traceback
import gc
import sys

# Configure garbage collection for more aggressive memory management
gc.enable()
gc.set_threshold(700, 10, 5)

# Set memory usage limit for streamlit (4GB)
if sys.platform.startswith('win'):
    import ctypes
    kernel32 = ctypes.windll.kernel32
    kernel32.SetProcessWorkingSetSize(-1, -1, -1)

def init_session_state():
    """Initialize session state variables"""
    try:
        if 'initialized' not in st.session_state:
            st.session_state.initialized = False
        
        if not st.session_state.initialized:
            with st.status("Initializing Learning Assistant...", expanded=True) as status:
                st.write("Loading documents and building knowledge base...")
                st.session_state.chatbot = LearningChatbot()
                st.session_state.messages = []
                st.session_state.error = None
                status.update(label="Ready!", state="complete", expanded=False)
            
            st.session_state.initialized = True
    except Exception as e:
        st.session_state.error = f"Initialization Error: {str(e)}"
        print(f"Error during initialization: {traceback.format_exc()}")

def display_error_message():
    """Display error message with a helpful tone"""
    st.error(
        "I'm having trouble getting started. Could you help me check:\n"
        "1. If the study materials are in the documents folder?\n"
        "2. If the API key is set up correctly?\n"
        "3. If all the required packages are installed?"
    )
    if st.session_state.error:
        with st.expander("Technical Details"):
            st.code(st.session_state.error)

def main():
    try:
        st.set_page_config(
            page_title="OpexA Career Assistant",
            page_icon="🚀",
            layout="wide",
            initial_sidebar_state="expanded"
        )
    except Exception as e:
        st.write("Error configuring page:", str(e))
        return

    # Initialize session
    init_session_state()

    # Check for errors
    if hasattr(st.session_state, 'error') and st.session_state.error:
        display_error_message()
        return

    # Header with OpexA branding
    col1, col2 = st.columns([3, 1])
    with col1:
        st.title("🚀 OpexA Career Assistant")
        st.caption("Your AI guide for IT career growth and skill development")
    
    # Enhanced sidebar with OpexA info
    with st.sidebar:
        st.info("""
        **Welcome to OpexA!** 🎯
        
        I'm your AI career assistant, here to help you:
        • Navigate your IT career path
        • Understand skill assessments (CMA & SLA)
        • Explore learning opportunities
        • Make informed career decisions
        
        Ask me anything about your professional journey!
        """)
        
        # Add subscription plan info
        with st.expander("💡 Available Plans"):
            st.markdown("""
            **Explorer (Free)**
            • Basic career guides
            • Initial assessments
            • Community access
            
            **Starter**
            • Advanced features
            • Full SLA access
            • Personalized learning paths
            • Premium resources
            """)
        
        # Clear conversation button with improved styling
        if st.button("🔄 Start Fresh"):
            st.session_state.messages = []
            st.rerun()

    # Chat interface
    try:
        # Display chat messages with improved formatting
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        # Enhanced chat input
        prompt = st.chat_input("Ask about your IT career journey, skills, or OpexA features...")

        if prompt:
            # Validate chatbot
            if not hasattr(st.session_state, 'chatbot') or not st.session_state.chatbot.verify_knowledge_base():
                st.error("Give me a moment to get ready. Please check if everything is set up correctly.")
                return

            # Show user message
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)

            # Generate and show response with improved error handling
            with st.chat_message("assistant"):
                with st.spinner("Analyzing your question..."):
                    try:
                        response = st.session_state.chatbot.generate_response(prompt)
                        st.markdown(response)
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": response
                        })
                    except Exception as e:
                        st.error("I apologize, but I'm having trouble understanding that. Could you rephrase your question?")
                        print(f"Error generating response: {str(e)}")

    except Exception as e:
        st.error("Something unexpected happened. Let's start fresh!")
        print(f"Error in main loop: {traceback.format_exc()}")

if __name__ == "__main__":
    main()