app.py

import streamlit as st
import time
import os
import sys
import json
from datetime import datetime
from pathlib import Path
sys.path.append(str(Path(__file__).parent))

# Import modules
from src.ui.chat_handler import chat_handler
from utils.config import config
from core.session import session_manager
from core.memory import check_redis_health
from core.errors import translate_error
from core.personality import personality
from src.analytics.user_logger import user_logger
from src.analytics.session_analytics import session_analytics
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

st.set_page_config(page_title="CosmicCat AI Assistant", page_icon="🐱", layout="wide")

# Initialize session state properly with error handling
def initialize_session_state():
    """Initialize all session state variables safely"""
    session_vars = {
        "messages": [],
        "is_processing": False,
        "ngrok_url_temp": st.session_state.get("ngrok_url", "https://7bcc180dffd1.ngrok-free.app"),
        "cosmic_mode": True,
        "show_welcome": True,
        "last_processed_message": "",
        "selected_model_value": "auto",
        "session_id": f"sess_{int(time.time())}_{abs(hash(str(time.time()))) % 10000}"
    }
    
    for var_name, default_value in session_vars.items():
        if var_name not in st.session_state:
            st.session_state[var_name] = default_value
            logger.info(f"Initialized session variable: {var_name}")

# Initialize session state
try:
    initialize_session_state()
except Exception as e:
    logger.error(f"Session state initialization failed: {e}")
    # Fallback initialization
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "is_processing" not in st.session_state:
        st.session_state.is_processing = False
    if "ngrok_url_temp" not in st.session_state:
        st.session_state.ngrok_url_temp = st.session_state.get("ngrok_url", "https://7bcc180dffd1.ngrok-free.app")
    if "cosmic_mode" not in st.session_state:
        st.session_state.cosmic_mode = True
    if "show_welcome" not in st.session_state:
        st.session_state.show_welcome = True
    if "last_processed_message" not in st.session_state:
        st.session_state.last_processed_message = ""
    if "selected_model_value" not in st.session_state:
        st.session_state.selected_model_value = "auto"
    if "session_id" not in st.session_state:
        st.session_state.session_id = f"sess_{int(time.time())}_{abs(hash(str(time.time()))) % 10000}"

# Start session tracking with error handling
try:
    session_analytics.start_session_tracking("default_user", st.session_state.session_id)
except Exception as e:
    logger.warning(f"Analytics session tracking failed: {e}")

# Log page view
session_analytics.track_interaction("default_user", st.session_state.session_id, "page_view", {
    "page": "main_chat"
})

# Sidebar
with st.sidebar:
    st.title("🐱 CosmicCat AI Assistant")
    st.markdown("Your personal AI-powered assistant with a cosmic twist.")
    
    # Model selection
    model_options = {
        "Auto Select": "auto",  # Will choose best available
        "🦙 Ollama (Local)": "ollama",
        "🤗 HF Endpoint": "huggingface"
    }
    
    selected_model_key = st.selectbox(
        "Select Provider",
        options=list(model_options.keys()),
        index=0
    )
    st.session_state.selected_model_value = model_options[selected_model_key]
    
    # Show which provider will actually be used
    actual_provider = "Unknown"
    if st.session_state.selected_model_value == "auto":
        if config.hf_token and hf_monitor.get_endpoint_status()["available"]:
            actual_provider = "🤗 HF Endpoint"
        elif config.ollama_host:
            actual_provider = "🦙 Ollama"
    else:
        actual_provider = "🤗 HF Endpoint" if st.session_state.selected_model_value == "huggingface" else "🦙 Ollama"
    
    st.info(f"**Using Provider:** {actual_provider}")
    
    # Log model selection
    session_analytics.track_interaction("default_user", st.session_state.session_id, "model_selection", {
        "selected_model": st.session_state.selected_model_value,
        "actual_provider": actual_provider
    })
    
    # Cosmic mode toggle
    st.session_state.cosmic_mode = st.checkbox("Enable Cosmic Mode", value=st.session_state.cosmic_mode,
                                              on_change=lambda: session_analytics.track_interaction("default_user", st.session_state.session_id, "cosmic_mode_toggle", {
                                                  "enabled": st.session_state.cosmic_mode
                                              }))
    
    st.divider()
    
    # Configuration
    st.subheader("⚙️ Configuration")
    ngrok_url_input = st.text_input(
        "Ollama Server URL",
        value=st.session_state.ngrok_url_temp,
        help="Enter your ngrok URL",
        on_change=lambda: session_analytics.track_interaction("default_user", st.session_state.session_id, "url_update", {
            "url_changed": ngrok_url_input != st.session_state.ngrok_url_temp
        })
    )
    
    if ngrok_url_input != st.session_state.ngrok_url_temp:
        st.session_state.ngrok_url_temp = ngrok_url_input
        st.success("✅ URL updated!")
        session_analytics.track_interaction("default_user", st.session_state.session_id, "url_updated", {
            "new_url": ngrok_url_input
        })
    
    if st.button("📡 Test Connection"):
        start_time = time.time()
        session_analytics.track_interaction("default_user", st.session_state.session_id, "test_connection_click")
        try:
            from core.providers.ollama import OllamaProvider
            ollama_provider = OllamaProvider(st.session_state.selected_model_value)
            is_valid = ollama_provider.validate_model()
            end_time = time.time()
            
            if is_valid:
                st.success("✅ Connection successful!")
                session_analytics.track_interaction("default_user", st.session_state.session_id, "connection_success", {
                    "response_time": end_time - start_time
                })
                user_logger.log_performance_metric("default_user", "connection_test", end_time - start_time)
            else:
                st.error("❌ Model validation failed")
                session_analytics.track_interaction("default_user", st.session_state.session_id, "connection_failed", {
                    "error": "model_validation_failed"
                })
        except Exception as e:
            end_time = time.time()
            st.error(f"❌ Error: {str(e)[:50]}...")
            session_analytics.track_interaction("default_user", st.session_state.session_id, "connection_error", {
                "error": str(e)[:100],
                "response_time": end_time - start_time
            })
            user_logger.log_error("default_user", "connection_test", str(e))
    
    if st.button("🗑️ Clear History"):
        session_analytics.track_interaction("default_user", st.session_state.session_id, "clear_history_click")
        st.session_state.messages = []
        st.session_state.last_processed_message = ""
        # Also clear backend session
        session_manager.clear_session("default_user")
        st.success("History cleared!")
        session_analytics.track_interaction("default_user", st.session_state.session_id, "history_cleared")
    
    st.divider()
    
    # System Status with enhanced HF monitoring
    with st.expander("🔍 System Status", expanded=True):
        st.subheader("📊 Status")
        
        # Ollama Status
        try:
            from services.ollama_monitor import check_ollama_status
            ollama_status = check_ollama_status()
            if ollama_status.get("running"):
                st.success("🦙 Ollama: Running")
            else:
                st.warning("🦙 Ollama: Not running")
        except:
            st.info("🦙 Ollama: Unknown")
        
        # HF Endpoint Status (Enhanced with initialization info)
        try:
            from src.services.hf_monitor import hf_monitor
            status_message = hf_monitor.get_human_readable_status()
            
            # Display appropriate status icon
            if "🟢" in status_message:
                st.success(status_message)
            elif "🟡" in status_message:
                st.warning(status_message)
            elif "🔴" in status_message:
                st.error(status_message)
            elif "❌" in status_message:
                st.error(status_message)
            elif "⏳" in status_message:
                st.info(status_message)
            else:
                st.info(status_message)
                
            # Show initialization progress if applicable
            init_progress = hf_monitor.get_initialization_progress()
            if init_progress:
                st.info(init_progress)
                
            # Add wake-up button if scaled to zero or initializing
            if "scaled to zero" in status_message.lower() or "initializing" in status_message.lower():
                if st.button("⚡ Wake Up HF Endpoint", key="wake_up_hf"):
                    session_analytics.track_interaction("default_user", st.session_state.session_id, "wake_up_hf_click")
                    with st.spinner("Attempting to wake up HF endpoint... This may take 2-4 minutes during initialization..."):
                        start_time = time.time()
                        if hf_monitor.attempt_wake_up():
                            end_time = time.time()
                            st.success("✅ Wake-up request sent! The endpoint should be initializing now. Try your request again in a moment.")
                            session_analytics.track_interaction("default_user", st.session_state.session_id, "hf_wake_up_success", {
                                "response_time": end_time - start_time
                            })
                            user_logger.log_performance_metric("default_user", "hf_wake_up", end_time - start_time)
                            time.sleep(3)
                            st.experimental_rerun()
                        else:
                            end_time = time.time()
                            st.error("❌ Failed to send wake-up request. Please try again or wait for initialization to complete.")
                            session_analytics.track_interaction("default_user", st.session_state.session_id, "hf_wake_up_failed", {
                                "response_time": end_time - start_time
                            })
                            user_logger.log_error("default_user", "hf_wake_up", "Failed to wake up HF endpoint")
                            
        except Exception as e:
            st.info(f"🤗 HF Endpoint: Error checking status - {str(e)}")
            session_analytics.track_interaction("default_user", st.session_state.session_id, "hf_status_error", {
                "error": str(e)
            })
            user_logger.log_error("default_user", "hf_status_check", str(e))
        
        # Redis Status
        try:
            if check_redis_health():
                st.success("💾 Redis: Connected")
            else:
                st.error("💾 Redis: Disconnected")
        except:
            st.info("💾 Redis: Unknown")
    
    st.divider()
    
    # Feedback Section
    st.subheader("⭐ Feedback")
    rating = st.radio("How would you rate your experience?", [1, 2, 3, 4, 5], horizontal=True)
    feedback_comment = st.text_area("Additional comments (optional):")
    if st.button("Submit Feedback"):
        user_logger.log_feedback("default_user", rating, feedback_comment)
        session_analytics.track_interaction("default_user", st.session_state.session_id, "feedback_submitted", {
            "rating": rating,
            "has_comment": bool(feedback_comment)
        })
        st.success("Thank you for your feedback! 🙏")
    
    st.divider()
    
    # Debug Info
    st.subheader("🐛 Debug Info")
    st.markdown(f"**Environment:** {'HF Space' if config.is_hf_space else 'Local'}")
    st.markdown(f"**Model:** {st.session_state.selected_model_value}")
    st.markdown(f"**Session ID:** {st.session_state.session_id}")

# Main interface
st.title("🐱 CosmicCat AI Assistant")
st.markdown("Ask me anything!")

# Welcome message
if st.session_state.show_welcome:
    with st.chat_message("assistant"):
        greeting = personality.get_greeting(cosmic_mode=st.session_state.cosmic_mode)
        st.markdown(greeting)
    st.session_state.show_welcome = False

# Display existing conversation history
for message in st.session_state.get("messages", []):
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        if "timestamp" in message:
            provider_info = f" (via {message.get('provider', 'ollama')})" if message["role"] == "assistant" else ""
            st.caption(f"🕒 {message['timestamp']}{provider_info}")

# Chat input with enhanced processing
user_input = st.chat_input("Type your message here...", key="chat_input")

# Process message when received
if user_input and user_input.strip():
    if not st.session_state.get('is_processing', False):
        chat_handler.process_user_message(user_input, st.session_state.selected_model_value)
    else:
        st.warning("Still processing your previous request...")

# Handle AI response processing (triggered after user message display)
if st.session_state.get('is_processing', False) and st.session_state.get('last_processed_message'):
    chat_handler.process_ai_response(
        st.session_state.last_processed_message, 
        st.session_state.selected_model_value
    )

# About tab
st.divider()
tab1, = st.tabs(["ℹ️ About"])

with tab1:
    st.header("ℹ️ About CosmicCat AI Assistant")
    st.markdown("""
    The CosmicCat AI Assistant is a sophisticated conversational AI with a cosmic theme.
    
    ### 🧠 Core Features
    - **Local AI processing** with Ollama models
    - **Persistent memory** using Redis
    - **Space-themed personality** for fun interactions
    - **HF Endpoint integration** for advanced capabilities
    
    ### 🚀 Cosmic Mode
    When enabled, the AI responds with space-themed language and metaphors.
    
    ### 🛠️ Technical Architecture
    - **Primary model**: HF Endpoint (advanced processing)
    - **Secondary model**: Ollama (local processing)
    - **Memory system**: Redis-based session management
    """)
    
    # Log about page view
    session_analytics.track_interaction("default_user", st.session_state.session_id, "about_page_view")

# End session tracking when app closes
def on_session_end():
    session_analytics.end_session_tracking("default_user", st.session_state.session_id)

# Register cleanup function
import atexit
atexit.register(on_session_end)