src/streamlit_app.py

import streamlit as st
import requests
import json
import os
from datetime import datetime
import time 
# For potential small delays to improve UI experience

# Import Firebase Admin SDK
import firebase_admin
from firebase_admin import credentials, firestore, auth
from firebase_admin.exceptions import FirebaseError

# --- Constants & Environment Variables ---
APP_ID = os.getenv('__app_id', 'default-app-id')
FIREBASE_CONFIG_STR = os.getenv('__firebase_config', '{}')
INITIAL_AUTH_TOKEN = os.getenv('__initial_auth_token', None)

OLLAMA_API_BASE_URL = os.getenv("OLLAMA_HOST", "http://localhost:11434")
OLLAMA_MODEL_NAME = 'krishna_choudhary/AI_Assistant_Chatbot'
# Your specific model name

# --- Firebase Initialization ---
if not firebase_admin._apps:
    try:
        # Use ApplicationDefault credentials which Canvas environment should provide
        cred = credentials.ApplicationDefault()
        firebase_admin.initialize_app(cred, name=APP_ID)
        st.session_state.db = firestore.client()
        st.session_state.auth = auth
        st.session_state.firebase_initialized = True
        print("Firebase Admin SDK initialized successfully.")
    except Exception as e:
        st.error(f"Error initializing Firebase Admin SDK: {e}")
        st.session_state.firebase_initialized = False
        print(f"Failed to initialize Firebase Admin SDK: {e}")

# --- Session State Initialization ---
if 'current_page' not in st.session_state:
    st.session_state.current_page = 'dashboard'
if 'messages' not in st.session_state:
    st.session_state.messages = []
if 'input_message' not in st.session_state:
    st.session_state.input_message = ''
if 'is_loading' not in st.session_state:
    st.session_state.is_loading = False
if 'current_chat_id' not in st.session_state:
    st.session_state.current_chat_id = None
if 'ollama_api_error' not in st.session_state:
    st.session_state.ollama_api_error = None
if 'user_id' not in st.session_state:
    st.session_state.user_id = None # Will be set after login/signup
if 'user_email' not in st.session_state:
    st.session_state.user_email = None # Store email for display/auth
if 'user_name' not in st.session_state:
    st.session_state.user_name = ""
if 'loaded_user_name' not in st.session_state:
    st.session_state.loaded_user_name = False
if 'raw_ollama_response' not in st.session_state:
    st.session_state.raw_ollama_response = None # To store full raw response for debug

# --- Custom CSS for "Code Mentor AI" Styling ---
st.markdown(
    """
    <style>
        /* General Styles */
        body {
            color: #E0E0E0;
            background-color: #0F172A;
        }
        .stApp {
            background-color: #0F172A; /* Main app background */
        }

        /* Main Content Area */
        .main .block-container {
            padding-top: 2rem;
            padding-bottom: 2rem;
            padding-left: 5rem;
            padding-right: 5rem;
        }

        /* Sidebar Styles */
        [data-testid="stSidebar"] {
            background-color: #1E293B;
            border-right: 1px solid #334155;
            padding-top: 2rem;
        }
        [data-testid="stSidebar"] h2 {
            color: #FFFFFF;
            font-weight: bold;
            text-align: center;
            margin-bottom: 2rem;
        }
        /* Sidebar buttons - selected/active state */
        [data-testid="stSidebar"] .stButton>button {
            background-color: transparent;
            color: #94A3B8;
            border: none;
            width: 100%;
            text-align: left;
            padding: 0.75rem 1rem;
            border-radius: 0.5rem;
            font-size: 1rem;
            display: flex;
            align-items: center;
            gap: 10px; /* Space between icon and text */
        }
        [data-testid="stSidebar"] .stButton>button:hover {
            background-color: #334155;
            color: #FFFFFF;
        }
        /* Streamlit applies an internal class to selected buttons.
           This targets the active button in the sidebar based on session state. */
        .sidebar-button-active {
            background-color: #4A5568 !important; /* Brighter background for active */
            color: #FFFFFF !important;
            font-weight: bold;
        }
        
        /* Chat Interface Styles */
        div[data-testid="stVerticalBlock"] > div[data-testid="stVerticalBlock"] > div[data-testid="stVerticalBlock"]:has([data-testid="stChatInput"]) {
            background-color: #1E293B;
            border-radius: 0.75rem;
            padding: 1rem;
            border-top: 1px solid #334155;
            margin-top: 1rem;
        }
        [data-testid="stChatInput"] {
            background-color: transparent !important;
        }
        [data-testid="stChatInput"] input {
            background-color: #2D3748 !important;
            color: #E0E0E0 !important;
            border: 1px solid #4A5568 !important;
            border-radius: 0.5rem !important;
            padding: 0.75rem 1rem;
        }
        [data-testid="stChatMessage"] {
            background-color: #334155;
            border-radius: 0.5rem;
            padding: 1rem;
            margin-bottom: 1rem;
        }
        .st-chat-message-user {
            background-color: #2D3748 !important;
            border: 1px solid #4A5568;
        }
        .st-chat-message-assistant {
            background-color: #1E293B !important;
            border: 1px solid #334155;
        }
        .stCodeBlock {
            background-color: #0F172A !important;
            border: 1px solid #334155 !important;
            border-radius: 0.5rem;
            padding: 1rem;
            overflow-x: auto;
        }

        /* Custom Cards for Stats */
        .stat-card {
            background-color: #1E293B;
            border-radius: 0.75rem;
            padding: 1.5rem;
            text-align: center;
            border: 1px solid #334155;
        }
        .stat-card h3 {
            color: #94A3B8;
            font-size: 1rem;
            margin-bottom: 0.5rem;
        }
        .stat-card p {
            color: #FFFFFF;
            font-size: 2.5rem;
            font-weight: bold;
        }

        /* Popular Topics styling */
        .popular-topics-card {
            background-color: #1E293B;
            border-radius: 0.75rem;
            padding: 1.5rem;
            border: 1px solid #334155;
        }
        .popular-topics-card h3 {
            color: #FFFFFF;
            font-size: 1.2rem;
            margin-bottom: 1rem;
        }
        .popular-topics-card ul {
            list-style-type: none;
            padding-left: 0;
        }
        .popular-topics-card li {
            color: #94A3B8;
            padding: 0.25rem 0;
            font-size: 0.95rem;
        }
        .popular-topics-card li::before {
            content: "❯ ";
            color: #69BFF8;
            font-weight: bold;
            margin-right: 0.5rem;
        }
        
        /* Info Box for How to Use */
        .info-box {
            background-color: #2D3748;
            border-left: 5px solid #69BFF8;
            border-radius: 0.5rem;
            padding: 1rem;
            margin-bottom: 1rem;
            color: #E0E0E0;
        }
        .info-box strong {
            color: #FFFFFF;
        }

        /* Custom Button for "Ask Coding Instructor" */
        .stButton.coding-instructor-button button {
            background-color: #69BFF8;
            color: white;
            font-size: 1.1rem;
            padding: 0.8rem 2rem;
            border-radius: 0.75rem;
            width: 100%;
            margin-top: 1.5rem;
            border: none;
            transition: background-color 0.3s ease;
        }
        .stButton.coding-instructor-button button:hover {
            background-color: #4CAF50;
        }

        /* Dark mode toggle and profile button in header */
        .header-buttons {
            display: flex;
            justify-content: flex-end;
            gap: 10px;
            margin-bottom: 1rem;
        }
        .header-buttons .stButton>button {
            background-color: #1E293B;
            border: 1px solid #334155;
            color: #94A3B8;
            padding: 0.5rem 1rem;
            border-radius: 0.5rem;
        }
    </style>
    """,
    unsafe_allow_html=True
)


# --- Firebase Helpers ---
def get_user_doc_ref(db_client, user_id):
    return db_client.collection(f"artifacts/{APP_ID}/users").document(user_id)

def get_chats_collection_ref(db_client, user_id):
    return get_user_doc_ref(db_client, user_id).collection("chats")

def get_settings_doc_ref(db_client, user_id):
    return get_user_doc_ref(db_client, user_id).collection("settings").document("userSettings")

# --- Authentication Functions ---
def login_user(username, password):
    if not st.session_state.firebase_initialized:
        st.error("Firebase not initialized.")
        return False

    email = f"{username}@app.com" # Map username to email for Firebase Auth
    try:
        user = st.session_state.auth.get_user_by_email(email)
        # In a real app, you'd verify the password here securely.
        # Firebase Admin SDK doesn't directly verify passwords from client.
        # This is a SIMULATION for Canvas environment.
        # A real client-side login would use signInWithEmailAndPassword.
        
        # For this server-side simulation, if user exists, we assume success.
        # This is INSECURE for production.
        st.session_state.user_id = user.uid
        st.session_state.user_email = user.email
        st.session_state.user_name = username
        st.success(f"Logged in as {username}!")
        st.rerun() # Rerun to update UI
        return True
    except FirebaseError as e:
        if "user-not-found" in str(e):
            st.error("User not found. Please sign up or check your username.")
        else:
            st.error(f"Login failed: {e}")
        return False
    except Exception as e:
        st.error(f"An unexpected error occurred during login: {e}")
        return False

def signup_user(username, password):
    if not st.session_state.firebase_initialized:
        st.error("Firebase not initialized.")
        return False

    email = f"{username}@app.com" # Map username to email for Firebase Auth
    try:
        user = st.session_state.auth.create_user(email=email, password=password)
        st.session_state.user_id = user.uid
        st.session_state.user_email = user.email
        st.session_state.user_name = username
        
        # Store username in Firestore settings for display
        settings_doc_ref = get_settings_doc_ref(st.session_state.db, user.uid)
        settings_doc_ref.set({'userName': username}, merge=True)

        st.success(f"Signed up and logged in as {username}!")
        st.rerun() # Rerun to update UI
        return True
    except FirebaseError as e:
        if "email-already-exists" in str(e):
            st.error("Username already exists. Please log in.")
        elif "invalid-password" in str(e):
            st.error("Password must be at least 6 characters long.")
        else:
            st.error(f"Signup failed: {e}")
        return False
    except Exception as e:
        st.error(f"An unexpected error occurred during signup: {e}")
        return False

def logout_user():
    st.session_state.user_id = None
    st.session_state.user_email = None
    st.session_state.user_name = ""
    st.session_state.current_chat_id = None
    st.session_state.messages = []
    st.session_state.ollama_api_error = None
    st.session_state.raw_ollama_response = None
    st.success("Logged out successfully.")
    st.rerun()

# --- Functions for Chatbot Logic ---
def load_chat_messages(chat_id):
    if not st.session_state.firebase_initialized or not st.session_state.user_id or not chat_id:
        return []
    try:
        messages_ref = get_chats_collection_ref(st.session_state.db, st.session_state.user_id).document(chat_id).collection("messages")
        messages_docs = messages_ref.order_by('timestamp').get()
        loaded_messages = []
        for doc in messages_docs:
            data = doc.to_dict()
            if 'timestamp' in data and hasattr(data['timestamp'], 'to_datetime'):
                data['timestamp'] = data['timestamp'].to_datetime()
            loaded_messages.append({'id': doc.id, **data})
        return loaded_messages
    except Exception as e:
        st.error(f"Error loading chat messages: {e}")
        return []

def start_new_chat_session():
    if not st.session_state.firebase_initialized or not st.session_state.user_id:
        st.session_state.ollama_api_error = "Database not ready. Cannot start new chat."
        return

    if not st.session_state.is_loading:
        st.session_state.is_loading = True
        st.session_state.ollama_api_error = None
        st.session_state.input_message = ''
        st.session_state.messages = []

        try:
            new_chat_ref = get_chats_collection_ref(st.session_state.db, st.session_state.user_id).add({
                'createdAt': firestore.SERVER_TIMESTAMP,
                'title': "New Chat"
            })
            st.session_state.current_chat_id = new_chat_ref[1].id
            print(f"New chat session started: {st.session_state.current_chat_id}")
        except Exception as e:
            st.session_state.ollama_api_error = f"Failed to start new chat: {e}"
            print(f"Error starting new chat: {e}")
        finally:
            st.session_state.is_loading = False
            st.rerun()

def send_message():
    user_message_content = st.session_state.input_message.strip()

    if not user_message_content or st.session_state.is_loading:
        return

    if not st.session_state.firebase_initialized or not st.session_state.user_id:
        st.session_state.ollama_api_error = "Database not ready. Cannot send message. Please log in."
        return

    st.session_state.is_loading = True
    st.session_state.ollama_api_error = None

    if not st.session_state.current_chat_id:
        st.session_state.is_loading = False
        start_new_chat_session()
        return

    chat_messages_ref = get_chats_collection_ref(st.session_state.db, st.session_state.user_id).document(st.session_state.current_chat_id).collection("messages")

    user_message_data = {
        'role': 'user',
        'content': user_message_content,
        'timestamp': firestore.SERVER_TIMESTAMP,
    }
    try:
        chat_messages_ref.add(user_message_data)
        st.session_state.messages.append({'id': 'temp_user', 'role': 'user', 'content': user_message_content, 'timestamp': datetime.now()})
        st.session_state.input_message = ''
    except Exception as e:
        st.session_state.ollama_api_error = f"Failed to save user message: {e}"
        st.session_state.is_loading = False
        st.rerun()
        return

    if st.session_state.current_chat_id and st.session_state.messages and len(st.session_state.messages) == 1:
        try:
            get_chats_collection_ref(st.session_state.db, st.session_state.user_id).document(st.session_state.current_chat_id).update({
                'title': user_message_content[:50] + "..." if len(user_message_content) > 50 else user_message_content
            })
        except Exception as e:
            print(f"Warning: Could not update new chat title: {e}")

    ollama_messages = [{'role': msg['role'], 'content': msg['content']} for msg in st.session_state.messages]

    try:
        response = requests.post(
            f'{OLLAMA_API_BASE_URL}/api/chat',
            json={
                'model': OLLAMA_MODEL_NAME,
                'messages': ollama_messages,
                'stream': False,
                'options': {
                    'temperature': 0.7,
                    'num_predict': 2048, # CRUCIAL: Increased max tokens for full responses
                    'top_p': 0.9,
                },
            },
            timeout=120
        )
        response.raise_for_status()
        result = response.json()
        st.session_state.raw_ollama_response = json.dumps(result, indent=2) # Store raw response

        assistant_response_content = result.get('message', {}).get('content', "No response from model.")

        chat_messages_ref.add({
            'role': 'assistant',
            'content': assistant_response_content,
            'timestamp': firestore.SERVER_TIMESTAMP,
        })
        st.session_state.messages.append({'id': 'temp_assistant', 'role': 'assistant', 'content': assistant_response_content, 'timestamp': datetime.now()})

    except requests.exceptions.Timeout:
        st.session_state.ollama_api_error = "Ollama connection timed out."
        print("Ollama connection timed out.")
    except requests.exceptions.RequestException as e:
        st.session_state.ollama_api_error = f"Ollama connection error: {e}. Is the Ollama server running and accessible?"
        print(f"Ollama connection error: {e}")
    except json.JSONDecodeError:
        st.session_state.ollama_api_error = "Ollama returned an unreadable response. Check server logs."
        print(f"Ollama JSON decode error: {response.text if 'response' in locals() else 'No response'}")
    except Exception as e:
        st.session_state.ollama_api_error = f"An unexpected error occurred: {e}"
        print(f"General Ollama error: {e}")
    finally:
        st.session_state.is_loading = False
        st.rerun()

# --- UI Page Components ---

def auth_ui():
    st.title("Welcome to Code Mentor AI")
    st.markdown("---")
    st.warning("""
    **Security Warning:** For demonstration purposes, this app uses a simplified username/password login.
    In a production environment, user-facing authentication should be handled by Firebase Client SDK
    or a secure backend API, not directly via Admin SDK in the frontend.
    """)

    auth_tab = st.tabs(["Login", "Sign Up"])

    with auth_tab[0]:
        st.subheader("Login")
        with st.form("login_form"):
            login_username = st.text_input("Username", key="login_username_input")
            login_password = st.text_input("Password", type="password", key="login_password_input")
            login_submitted = st.form_submit_button("Login")

            if login_submitted:
                if login_username and login_password:
                    login_user(login_username, login_password)
                else:
                    st.error("Please enter both username and password.")

    with auth_tab[1]:
        st.subheader("Sign Up")
        with st.form("signup_form"):
            signup_username = st.text_input("New Username", key="signup_username_input")
            signup_password = st.text_input("New Password", type="password", key="signup_password_input")
            signup_submitted = st.form_submit_button("Sign Up")

            if signup_submitted:
                if signup_username and signup_password:
                    signup_user(signup_username, signup_password)
                else:
                    st.error("Please enter both username and password.")

def dashboard_ui():
    # Header buttons (Dark Mode, Profile) - Placeholder functionality
    header_cols = st.columns([0.8, 0.1, 0.1])
    with header_cols[1]:
        st.button("🌙", key="dark_mode_btn", help="Toggle Dark Mode")
    with header_cols[2]:
        if st.button("👤", key="profile_btn", help="User Profile"):
            st.session_state.current_page = 'settings' # Go to settings on profile click
            st.rerun()

    st.title("Coding Instructor AI")
    st.markdown("---")

    # --- Popular Topics ---
    st.markdown("""
    <div class="popular-topics-card">
        <h3><span style="color: #F8B400; font-size: 1.2rem;">💡</span> Popular Topics</h3>
        <ul>
            <li>JavaScript Closures</li>
            <li>Python Decorators</li>
            <li>React Hooks</li>
            <li>Recursion Patterns</li>
            <li>Async/Await</li>
        </ul>
    </div>
    """, unsafe_allow_html=True)
    
    st.markdown("<br>", unsafe_allow_html=True)

    # --- Stats Cards ---
    col1, col2 = st.columns(2)
    with col1:
        st.markdown('<div class="stat-card"><h3>Questions Solved</h3><p>1,248</p></div>', unsafe_allow_html=True)
    with col2:
        st.markdown('<div class="stat-card"><h3>Languages</h3><p>24</p></div>', unsafe_allow_html=True)
    
    st.markdown("<br>", unsafe_allow_html=True)

    # --- Ask a Coding Question Section ---
    st.subheader("💬 Ask a Coding Question")

    # How to use info box
    st.markdown("""
    <div class="info-box">
        <strong>ⓘ How to use:</strong> Ask any coding-related question in any programming language. The AI is specialized to help with coding problems and concepts.
        <br>For non-coding questions, responses may be unpredictable.
    </div>
    """, unsafe_allow_html=True)

    # Initial load of messages for the current chat ID
    if st.session_state.current_chat_id and not st.session_state.messages:
        with st.spinner("Loading chat..."):
            st.session_state.messages = load_chat_messages(st.session_state.current_chat_id)
            time.sleep(0.5)

    # Display messages
    chat_container = st.container(height=400, border=False)
    with chat_container:
        if not st.session_state.messages and not st.session_state.is_loading:
            st.session_state.messages.append({"role": "assistant", "content": "How can I help you with your coding questions today?", "timestamp": datetime.now()})
            # No rerun here, let the main loop handle it to avoid infinite reruns on first load
        
        for msg in st.session_state.messages:
            message_class = "st-chat-message-user" if msg['role'] == 'user' else "st-chat-message-assistant"
            st.markdown(f'<div class="stChatMessage {message_class}">', unsafe_allow_html=True)
            st.markdown(msg['content'])
            if 'timestamp' in msg and msg['timestamp']:
                st.caption(f"Sent: {msg['timestamp'].strftime('%Y-%m-%d %H:%M:%S')}")
            st.markdown('</div>', unsafe_allow_html=True)

    if st.session_state.ollama_api_error:
        st.error(st.session_state.ollama_api_error)

    with st.form("chat_form", clear_on_submit=True):
        user_input = st.text_input(
            "Your Coding Question",
            key="chat_input_text",
            placeholder="e.g., Explain closures in JavaScript. How to implement binary search in Python. What is recursion?",
            disabled=st.session_state.is_loading or not st.session_state.user_id,
            label_visibility="visible"
        )
        
        col_send_new = st.columns([1, 1])
        with col_send_new[0]:
            send_button = st.form_submit_button(
                "Ask Coding Instructor",
                disabled=st.session_state.is_loading or not user_input.strip() or not st.session_state.user_id,
                help="Send your coding question to the AI.",
                use_container_width=True
            )
        with col_send_new[1]:
            new_chat_button = st.form_submit_button(
                "Start New Chat",
                disabled=st.session_state.is_loading or not st.session_state.user_id,
                help="Clear current chat and start a new conversation.",
                use_container_width=True
            )

        if send_button:
            st.session_state.input_message = user_input
            send_message()
        elif new_chat_button:
            start_new_chat_session()


def history_ui():
    st.title("📜 Chat History")
    st.markdown("---")
    if not st.session_state.firebase_initialized or not st.session_state.user_id:
        st.warning("Database not ready or user not logged in. Cannot load history.")
        return

    st.write(f"Logged in as: **{st.session_state.user_name}** (User ID: `{st.session_state.user_id}`)")

    try:
        chats_ref = get_chats_collection_ref(st.session_state.db, st.session_state.user_id)
        chat_docs = chats_ref.order_by('createdAt', direction=firestore.Query.DESCENDING).get()

        if not chat_docs:
            st.info("No chat sessions found. Start a new chat from the Dashboard!")
            return

        for chat_doc in chat_docs:
            chat_data = chat_doc.to_dict()
            chat_id = chat_doc.id
            created_at = chat_data.get('createdAt')
            if created_at and hasattr(created_at, 'strftime'):
                 created_at_str = created_at.strftime('%Y-%m-%d %H:%M')
            else:
                 created_at_str = "Unknown Date"

            title = chat_data.get('title', f"Chat ({created_at_str})")

            col_chat, col_delete = st.columns([5, 1])
            with col_chat:
                if st.button(f"**{title}**", key=f"chat_{chat_id}_select", use_container_width=True):
                    st.session_state.current_chat_id = chat_id
                    st.session_state.messages = load_chat_messages(chat_id)
                    st.session_state.current_page = 'dashboard'
                    st.rerun()
            with col_delete:
                if st.button("Delete", key=f"chat_{chat_id}_delete", use_container_width=True):
                    # Confirmation for deletion
                    if st.checkbox(f"Confirm delete '{title}'?", key=f"confirm_del_{chat_id}"):
                        delete_chat_and_messages(chat_id, chats_ref)
                        st.success(f"Chat '{title}' deleted.")
                        st.rerun()
    except Exception as e:
        st.error(f"Error fetching chat history: {e}")

def delete_chat_and_messages(chat_id, chats_ref):
    messages_ref = chats_ref.document(chat_id).collection("messages")
    for msg_doc in messages_ref.get():
        msg_doc.reference.delete()
    chats_ref.document(chat_id).delete()
    if st.session_state.current_chat_id == chat_id:
        st.session_state.current_chat_id = None
        st.session_state.messages = []


def tutorials_ui():
    st.title("📚 Tutorials")
    st.markdown("---")
    st.info("This section is under construction. Future updates will include interactive coding tutorials!")

def playground_ui():
    st.title("▶️ Playground")
    st.markdown("---")
    st.info("This section is under construction. Future updates will include a coding playground to test your code.")

def settings_ui():
    st.title("⚙️ Settings")
    st.markdown("---")
    if not st.session_state.firebase_initialized or not st.session_state.user_id:
        st.warning("Database not ready or user not logged in. Cannot load/save settings.")
        return

    settings_doc_ref = get_settings_doc_ref(st.session_state.db, st.session_state.user_id)

    if not st.session_state.loaded_user_name:
        try:
            settings_snap = settings_doc_ref.get()
            if settings_snap.exists:
                st.session_state.user_name = settings_snap.to_dict().get('userName', '')
            st.session_state.loaded_user_name = True
        except Exception as e:
            st.error(f"Error loading user settings: {e}")
            st.session_state.loaded_user_name = True

    new_user_name = st.text_input("Your Name:", value=st.session_state.user_name, key="settings_user_name_input")
    if st.button("Save Settings"):
        if new_user_name.strip():
            try:
                settings_doc_ref.set({'userName': new_user_name.strip()}, merge=True)
                st.session_state.user_name = new_user_name.strip()
                st.success("Settings saved successfully!")
            except Exception as e:
                st.error(f"Failed to save settings: {e}")
        else:
            st.warning("User name cannot be empty.")

    st.info(f"Your current User ID: `{st.session_state.user_id}`")
    if st.session_state.user_email:
        st.info(f"Your associated email: `{st.session_state.user_email}`")
    st.markdown("---")
    st.write("### Authentication Status")
    st.write("For persistent data across sessions, ensure your Hugging Face Space is configured with persistent authentication or link your account.")
    st.write("Currently, the app uses a session-based ID or the `__initial_auth_token` provided by the Canvas environment.")
    
    st.button("Logout", on_click=logout_user)

    st.markdown("---")
    st.write("### Raw Ollama Response (Debug)")
    if st.session_state.raw_ollama_response:
        st.code(st.session_state.raw_ollama_response, language="json")
    else:
        st.info("No raw Ollama response available yet. Send a message in the chat to see it here.")


# --- Main App Navigation and Page Rendering ---
# Sidebar navigation
with st.sidebar:
    st.markdown("## 🤖 Code Mentor AI")
    
    if st.session_state.user_id: # Show user name if logged in
        st.markdown(f"<p style='text-align: center; color: #69BFF8; font-weight: bold;'>Welcome, {st.session_state.user_name}!</p>", unsafe_allow_html=True)
    
    st.markdown("<br>", unsafe_allow_html=True)

    def set_page_and_rerun(page_name):
        st.session_state.current_page = page_name
        st.rerun()

    # Dashboard button
    dashboard_button_class = "sidebar-button-active" if st.session_state.current_page == "dashboard" else ""
    st.markdown(f'<button class="stButton secondary-button {dashboard_button_class}" onclick="window.parent.document.querySelector(\'[data-testid=stSidebar] button[key=sidebar_dashboard]\').click()" key="sidebar_dashboard">💬 Dashboard</button>', unsafe_allow_html=True)
    st.button("Dashboard", on_click=set_page_and_rerun, args=("dashboard",), key="sidebar_dashboard",
              type="secondary", use_container_width=True,
              help="Go to the main dashboard")
    
    # History button
    history_button_class = "sidebar-button-active" if st.session_state.current_page == "history" else ""
    st.markdown(f'<button class="stButton secondary-button {history_button_class}" onclick="window.parent.document.querySelector(\'[data-testid=stSidebar] button[key=sidebar_history]\').click()" key="sidebar_history">📚 History</button>', unsafe_allow_html=True)
    st.button("History", on_click=set_page_and_rerun, args=("history",), key="sidebar_history",
              type="secondary", use_container_width=True,
              help="View your past conversations")
    
    # Tutorials button
    tutorials_button_class = "sidebar-button-active" if st.session_state.current_page == "tutorials" else ""
    st.markdown(f'<button class="stButton secondary-button {tutorials_button_class}" onclick="window.parent.document.querySelector(\'[data-testid=stSidebar] button[key=sidebar_tutorials]\').click()" key="sidebar_tutorials">📖 Tutorials</button>', unsafe_allow_html=True)
    st.button("Tutorials", on_click=set_page_and_rerun, args=("tutorials",), key="sidebar_tutorials",
              type="secondary", use_container_width=True,
              help="Access coding tutorials")
              
    # Playground button
    playground_button_class = "sidebar-button-active" if st.session_state.current_page == "playground" else ""
    st.markdown(f'<button class="stButton secondary-button {playground_button_class}" onclick="window.parent.document.querySelector(\'[data-testid=stSidebar] button[key=sidebar_playground]\').click()" key="sidebar_playground">▶️ Playground</button>', unsafe_allow_html=True)
    st.button("Playground", on_click=set_page_and_rerun, args=("playground",), key="sidebar_playground",
              type="secondary", use_container_width=True,
              help="Experiment with code snippets")
              
    # Settings button
    settings_button_class = "sidebar-button-active" if st.session_state.current_page == "settings" else ""
    st.markdown(f'<button class="stButton secondary-button {settings_button_class}" onclick="window.parent.document.querySelector(\'[data-testid=stSidebar] button[key=sidebar_settings]\').click()" key="sidebar_settings">⚙️ Settings</button>', unsafe_allow_html=True)
    st.button("Settings", on_click=set_page_and_rerun, args=("settings",), key="sidebar_settings",
              type="secondary", use_container_width=True,
              help="Adjust application settings")

    st.markdown("""
        ---
        <div style="text-align:center; color: #94A3B8; font-size: 0.8rem;">
            Code Mentor AI v2.0<br>
            Powered by Streamlit & Ollama
        </div>
    """, unsafe_allow_html=True)


# Render Current Page based on st.session_state.current_page
if not st.session_state.user_id:
    auth_ui()
elif st.session_state.current_page == 'dashboard':
    dashboard_ui()
elif st.session_state.current_page == 'history':
    history_ui()
elif st.session_state.current_page == 'tutorials':
    tutorials_ui()
elif st.session_state.current_page == 'playground':
    playground_ui()
elif st.session_state.current_page == 'settings':
    settings_ui()

# Initial chat session creation on first run if no chat is selected and user is logged in
if (st.session_state.firebase_initialized and st.session_state.user_id and
    st.session_state.current_page == 'dashboard' and
    not st.session_state.current_chat_id and
    not st.session_state.messages):
    start_new_chat_session()