app.py

import streamlit as st
import time
import os
import sys
import json
import asyncio
from datetime import datetime
from pathlib import Path
sys.path.append(str(Path(__file__).parent))
from utils.config import config
from core.session import session_manager
from core.memory import check_redis_health
from core.coordinator import coordinator
from core.errors import translate_error
from core.personality import personality
from services.hf_endpoint_monitor import hf_monitor
from services.weather import weather_service
from core.llm import LLMClient
from core.providers.ollama import OllamaProvider
from core.providers.huggingface import HuggingFaceProvider
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 safely at the top of app.py
if "messages" not in st.session_state:
    st.session_state.messages = []
if "last_error" not in st.session_state:
    st.session_state.last_error = ""
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  # Default to cosmic mode
if "show_welcome" not in st.session_state:
    st.session_state.show_welcome = True

# Sidebar layout redesign
with st.sidebar:
    st.title("🐱 CosmicCat AI Assistant")
    st.markdown("Your personal AI-powered life development assistant")
    
    # PRIMARY ACTIONS
    st.subheader("💬 Primary Actions")
    model_options = {
        "Mistral 7B (Local)": "mistral:latest",
        "Llama 2 7B (Local)": "llama2:latest",
        "OpenChat 3.5 (Local)": "openchat:latest"
    }
    selected_model_name = st.selectbox(
        "Select Model",
        options=list(model_options.keys()),
        index=0,
        key="sidebar_model_select"
    )
    st.session_state.selected_model = model_options[selected_model_name]
    
    # Toggle for cosmic mode using checkbox
    st.session_state.cosmic_mode = st.checkbox("Enable Cosmic Mode", value=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",
        key="sidebar_ngrok_url"
    )
    
    if ngrok_url_input != st.session_state.ngrok_url_temp:
        st.session_state.ngrok_url_temp = ngrok_url_input
        st.success("✅ URL updated!")
        
    if st.button("📡 Test Connection"):
        try:
            # Use OllamaProvider to test connection
            ollama_provider = OllamaProvider(st.session_state.selected_model)
            # Test model validation
            is_valid = ollama_provider.validate_model()
            if is_valid:
                st.success("✅ Connection successful!")
            else:
                st.error("❌ Model validation failed")
        except Exception as e:
            st.error(f"❌ Error: {str(e)[:50]}...")
            
    if st.button("🗑️ Clear History"):
        st.session_state.messages = []
        st.success("History cleared!")
        
    st.divider()
    
    # SYSTEM STATUS
    with st.expander("🔍 System Status", expanded=False):
        st.subheader("📊 System Monitor")
        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")
            
        try:
            hf_status = hf_monitor.check_endpoint_status()
            # Enhanced HF status display with cat-themed messages
            if hf_status.get('available'):
                if hf_status.get('initialized', False):
                    st.success(f"🤗 HF Endpoint: Available ({hf_status.get('status_code')} OK)")
                    if hf_status.get('model'):
                        st.info(f"   Model: {hf_status.get('model')}")
                    if hf_status.get('region'):
                        st.info(f"   Region: {hf_status.get('region')}")
                    if hf_status.get('warmup_count'):
                        st.info(f"   Warmup Count: {hf_status.get('warmup_count')}")
                else:
                    st.warning("⏳ Kittens Waking Up...")
            elif hf_status.get('status_code') == 200:
                st.info("📡 Calling Space Friends...")
            else:
                st.error("😴 Nap Cat")
        except Exception as e:
            st.info("⏳ Kittens Stretching...")
            
        if check_redis_health():
            st.success("💾 Redis: Connected")
        else:
            st.error("💾 Redis: Disconnected")
            
    st.divider()
    
    st.subheader("🐛 Debug Info")
    # Show enhanced debug information
    st.markdown(f"**Environment:** {'HF Space' if config.is_hf_space else 'Local'}")
    st.markdown(f"**Model:** {st.session_state.selected_model}")
    st.markdown(f"**Fallback:** {'Enabled' if config.use_fallback else 'Disabled'}")
    
    # Show active features
    features = []
    if os.getenv("TAVILY_API_KEY"):
        features.append("Web Search")
    if config.openweather_api_key:
        features.append("Weather")
    st.markdown(f"**Active Features:** {', '.join(features) if features else 'None'}")
    
    # Show recent activity
    try:
        user_session = session_manager.get_session("default_user")
        coord_stats = user_session.get('ai_coordination', {})
        if coord_stats and coord_stats.get('last_coordination'):
            st.markdown(f"**Last Request:** {coord_stats.get('last_coordination')}")
        else:
            st.markdown("**Last Request:** N/A")
    except:
        st.markdown("**Last Request:** N/A")
        
    # Show Ollama ping status
    try:
        import requests
        import time
        start_time = time.time()
        headers = {
            "ngrok-skip-browser-warning": "true",
            "User-Agent": "CosmicCat-Debug"
        }
        response = requests.get(
            f"{st.session_state.ngrok_url_temp}/api/tags",
            headers=headers,
            timeout=15
        )
        ping_time = round((time.time() - start_time) * 1000)
        if response.status_code == 200:
            st.markdown(f"**Ollama Ping:** {response.status_code} OK ({ping_time}ms)")
        else:
            st.markdown(f"**Ollama Ping:** {response.status_code} Error")
    except Exception as e:
        st.markdown("**Ollama Ping:** Unreachable")
        
    # Redis status
    if check_redis_health():
        st.markdown("**Redis:** Healthy")
    else:
        st.markdown("**Redis:** Unhealthy")

# Main interface
st.title("🐱 CosmicCat AI Assistant")
st.markdown("Ask me anything about personal development, goal setting, or life advice!")

# Show welcome message only once
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

# Consistent message rendering function with cosmic styling
def render_message(role, content, source=None, timestamp=None):
    """Render chat messages with consistent styling"""
    with st.chat_message(role):
        if source:
            if source == "local_kitty":
                st.markdown(f"### 🐱 Cosmic Kitten Says:")
            elif source == "orbital_station":
                st.markdown(f"### 🛰️ Orbital Station Reports:")
            elif source == "cosmic_summary":
                st.markdown(f"### 🌟 Final Cosmic Summary:")
            elif source == "error":
                st.markdown(f"### ❌ Error:")
            elif source == "space_story":
                st.markdown(f"### 🐱 Cosmic Kitten Story:")
            else:
                st.markdown(f"### {source}")
        st.markdown(content)
        if timestamp:
            st.caption(f"🕒 {timestamp}")

# Display messages
for message in st.session_state.messages:
    render_message(
        message["role"],
        message["content"],
        message.get("source"),
        message.get("timestamp")
    )

# Input validation function
def validate_user_input(text):
    """Validate and sanitize user input"""
    if not text or not text.strip():
        return False, "Input cannot be empty"
    if len(text) > 1000:
        return False, "Input too long (max 1000 characters)"
    
    # Check for potentially harmful patterns
    harmful_patterns = ["<script", "javascript:", "onload=", "onerror="]
    if any(pattern in text.lower() for pattern in harmful_patterns):
        return False, "Potentially harmful input detected"
    
    return True, text.strip()

# Chat input - FIXED VERSION
user_input = st.chat_input("Type your message here...", disabled=st.session_state.is_processing)

# Process message when received
if user_input and not st.session_state.is_processing:
    # Validate input
    is_valid, validated_input = validate_user_input(user_input)
    if not is_valid:
        st.error(validated_input)
        st.session_state.is_processing = False
        st.experimental_rerun()  # Fixed: use experimental_rerun
    else:
        st.session_state.is_processing = True
        
        # Display user message immediately
        with st.chat_message("user"):
            st.markdown(validated_input)
            
        # Add to message history - ensure proper format
        st.session_state.messages.append({
            "role": "user",
            "content": validated_input,
            "timestamp": datetime.now().strftime("%H:%M:%S")
        })
        
        # Process AI response
        response_container = st.empty()
        status_placeholder = st.empty()
        response_placeholder = st.empty()
        
        try:
            # Get conversation history from session
            user_session = session_manager.get_session("default_user")
            conversation_history = user_session.get("conversation", []).copy()
            
            # Add the current user message to history for context
            conversation_history.append({"role": "user", "content": validated_input})
            
            # Check if cosmic mode is enabled
            if st.session_state.cosmic_mode:
                # Process cosmic cascade response
                status_placeholder.info("🐱 Cosmic Kitten Responding...")
                
                try:
                    # Get conversation history
                    user_session = session_manager.get_session("default_user")
                    conversation_history = user_session.get("conversation", []).copy()
                    conversation_history.append({"role": "user", "content": validated_input})
                    
                    # Stage 1: Local Ollama Response
                    ollama_provider = OllamaProvider(st.session_state.selected_model)
                    local_response = ollama_provider.generate(validated_input, conversation_history)
                    
                    if local_response:
                        # Display response (no nested st.chat_message)
                        st.markdown(f"### 🐱 Cosmic Kitten Says:\n{local_response}")
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": local_response,
                            "source": "local_kitty",
                            "timestamp": datetime.now().strftime("%H:%M:%S")
                        })
                    
                    # Stage 2: HF Endpoint Analysis
                    status_placeholder.info("🛰️ Beaming Query to Orbital Station...")
                    if config.hf_token:
                        # Check HF status first
                        hf_status = hf_monitor.check_endpoint_status()
                        if not hf_status['available']:
                            status_placeholder.info(personality.get_initializing_message())
                        
                        hf_provider = HuggingFaceProvider("meta-llama/Llama-2-7b-chat-hf")
                        hf_response = hf_provider.generate(validated_input, conversation_history)
                        
                        if hf_response:
                            # Display response (no nested st.chat_message)
                            st.markdown(f"### 🛰️ Orbital Station Reports:\n{hf_response}")
                            st.session_state.messages.append({
                                "role": "assistant",
                                "content": hf_response,
                                "source": "orbital_station",
                                "timestamp": datetime.now().strftime("%H:%M:%S")
                            })
                    
                    # Stage 3: Local Synthesis
                    status_placeholder.info("🐱 Cosmic Kitten Synthesizing Wisdom...")
                    
                    # Update history with both responses
                    synthesis_history = conversation_history.copy()
                    synthesis_history.extend([
                        {"role": "assistant", "content": local_response},
                        {"role": "assistant", "content": hf_response, "source": "cloud"}
                    ])
                    
                    synthesis = ollama_provider.generate(
                        f"Synthesize these two perspectives:\n1. Local: {local_response}\n2. Cloud: {hf_response}",
                        synthesis_history
                    )
                    
                    if synthesis:
                        # Display response (no nested st.chat_message)
                        st.markdown(f"### 🌟 Final Cosmic Summary:\n{synthesis}")
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": synthesis,
                            "source": "cosmic_summary",
                            "timestamp": datetime.now().strftime("%H:%M:%S")
                        })
                    
                    status_placeholder.success("✨ Cosmic Cascade Complete!")
                    
                except Exception as e:
                    error_msg = f"🌌 Cosmic disturbance: {str(e)}"
                    st.error(error_msg)
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": error_msg,
                        "source": "error",
                        "timestamp": datetime.now().strftime("%H:%M:%S")
                    })
            else:
                # Traditional processing
                # Try Ollama first
                status_placeholder.info("🦙 Contacting Ollama...")
                ai_response = None
                
                try:
                    # Use the OllamaProvider directly with proper configuration
                    ollama_provider = OllamaProvider(st.session_state.selected_model)
                    ai_response = ollama_provider.generate(validated_input, conversation_history)
                    
                    if ai_response:
                        st.markdown(ai_response)  # Use st.markdown instead of response_placeholder
                        status_placeholder.success("✅ Response received!")
                    else:
                        status_placeholder.warning("⚠️ Empty response from Ollama")
                        
                except Exception as ollama_error:
                    error_message = str(ollama_error)
                    status_placeholder.error(f"❌ Ollama error: {error_message[:100]}...")
                    logger.error(f"Ollama error: {error_message}")
                    
                    # Fallback to HF if available
                    if config.hf_token and not ai_response:
                        status_placeholder.info("⚡ Initializing HF Endpoint (2–4 minutes)...")
                        
                        try:
                            # Check HF status first
                            hf_status = hf_monitor.check_endpoint_status()
                            if not hf_status['available']:
                                status_placeholder.info(personality.get_initializing_message())
                                
                            # Use the HuggingFaceProvider directly
                            hf_provider = HuggingFaceProvider("meta-llama/Llama-2-7b-chat-hf")
                            ai_response = hf_provider.generate(validated_input, conversation_history)
                            
                            if ai_response:
                                st.markdown(ai_response)  # Use st.markdown instead of response_placeholder
                                status_placeholder.success("✅ HF response received!")
                            else:
                                status_placeholder.error("❌ No response from HF")
                                
                        except Exception as hf_error:
                            error_message = str(hf_error)
                            status_placeholder.error(f"❌ HF also failed: {error_message[:100]}...")
                            logger.error(f"HF error: {error_message}")
                            
                # Save response if successful
                if ai_response:
                    # Update conversation history in session
                    conversation = user_session.get("conversation", []).copy()
                    conversation.append({"role": "user", "content": validated_input})
                    conversation.append({"role": "assistant", "content": ai_response})
                    session_manager.update_session("default_user", {"conversation": conversation})
                    
                    # Add to message history
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": ai_response,
                        "timestamp": datetime.now().strftime("%H:%M:%S")
                    })
                else:
                    error_msg = "Sorry, I couldn't process your request. Please try again."
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": error_msg,
                        "timestamp": datetime.now().strftime("%H:%M:%S")
                    })
                    st.markdown(error_msg)
                    
        except Exception as e:
            error_msg = f"System error: {str(e)}"
            logger.error(f"Chat processing error: {error_msg}")
            st.error(error_msg)
            st.session_state.messages.append({
                "role": "assistant",
                "content": error_msg,
                "timestamp": datetime.now().strftime("%H:%M:%S")
            })
        finally:
            st.session_state.is_processing = False
            st.experimental_rerun()  # Fixed: use experimental_rerun

# Add evaluation dashboard tab (separate from chat interface) - ONLY ABOUT TAB NOW
st.divider()
# Only one tab now - About
tab1, = st.tabs(["ℹ️ About"])

with tab1:
    st.header("ℹ️ About CosmicCat AI Assistant")
    st.markdown("""
    The CosmicCat AI Assistant is a sophisticated conversational AI system with the following capabilities:
    
    ### 🧠 Core Features
    - **Multi-model coordination**: Combines local Ollama models with cloud-based Hugging Face endpoints
    - **Live web search**: Integrates with Tavily API for current information
    - **Persistent memory**: Uses Redis for conversation history storage
    - **Hierarchical reasoning**: Fast local responses with deep cloud analysis
    
    ### 🚀 Cosmic Mode
    When enabled, the AI follows a three-stage response pattern:
    1. **🐱 Cosmic Kitten Response**: Immediate local processing
    2. **🛰️ Orbital Station Analysis**: Deep cloud-based analysis
    3. **🌟 Final Synthesis**: Unified response combining both perspectives
    
    ### 🛠️ Technical Architecture
    - **Primary model**: Ollama (local processing for fast responses)
    - **Secondary model**: Hugging Face Inference API (deep analysis)
    - **External data**: Web search, weather data, and space information
    - **Memory system**: Redis-based session management
    
    ### 📊 Evaluation Tools
    - Behavior testing with sample prompts
    - Performance metrics and analytics
    """)

# Add special command handling for stories
if user_input and user_input.lower().strip() in ["tell me a story", "tell me a cosmic cat story", "story", "cosmic story", "tell me a space story"]:
    story = personality.get_space_story()
    st.markdown(f"### 🐱 Cosmic Kitten Story:\n\n{story}")
    st.session_state.messages.append({
        "role": "assistant",
        "content": story,
        "source": "space_story",
        "timestamp": datetime.now().strftime("%H:%M:%S")
    })
    st.session_state.is_processing = False
    st.experimental_rerun()