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.llm import send_to_ollama, send_to_hf
from core.session import session_manager
from core.memory import check_redis_health
from core.coordinator import coordinator
from core.errors import translate_error
import logging

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

st.set_page_config(page_title="AI Life Coach", page_icon="🧠", layout="wide")

# Processing stage labels for better user feedback
PROCESSING_STAGES = {
    "ollama": "🦙 Contacting Ollama...",
    "hf_init": "⚡ Initializing HF Endpoint (2–4 minutes)...",
    "hf_thinking": "🧠 HF Expert Thinking...",
    "hf_complete": "🎯 HF Analysis Complete!",
    "error": "⚠️ Something went wrong – trying again..."
}

# Initialize session state
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 "hf_expert_requested" not in st.session_state:
    st.session_state.hf_expert_requested = False

# Sidebar layout redesign
with st.sidebar:
    st.title("🧠 AI Life Coach")
    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]
    
    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:
            import requests
            headers = {
                "ngrok-skip-browser-warning": "true",
                "User-Agent": "AI-Life-Coach-Test"
            }
            with st.spinner("Testing connection..."):
                response = requests.get(
                    f"{ngrok_url_input}/api/tags",
                    headers=headers,
                    timeout=15
                )
                if response.status_code == 200:
                    st.success("✅ Connection successful!")
                else:
                    st.error(f"❌ Failed: {response.status_code}")
        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()
    
    # ADVANCED FEATURES
    with st.expander("🔍 Advanced Features", 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:
            from services.hf_endpoint_monitor import hf_monitor
            hf_status = hf_monitor.check_endpoint_status()
            if hf_status['available']:
                st.success("🤗 HF: Available")
            else:
                st.warning("🤗 HF: Not available")
        except:
            st.info("🤗 HF: Unknown")
        
        if check_redis_health():
            st.success("💾 Redis: Connected")
        else:
            st.error("💾 Redis: Disconnected")
        
        st.divider()
        
        st.subheader("🤖 HF Expert Analysis")
        col1, col2 = st.columns([3, 1])
        with col1:
            st.markdown("""
                **HF Expert Features:**
                - Analyzes entire conversation history
                - Performs web research when needed
                - Provides deep insights and recommendations
                - Acts as expert consultant in your conversation
            """)
        with col2:
            if st.button("🧠 Activate HF Expert",
                         key="activate_hf_expert_sidebar",
                         help="Send conversation to HF endpoint for deep analysis",
                         use_container_width=True,
                         disabled=st.session_state.is_processing):
                st.session_state.hf_expert_requested = True

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

# Display messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        # Format HF expert messages differently
        if message.get("source") == "hf_expert":
            st.markdown("### 🤖 HF Expert Analysis")
            st.markdown(message["content"])
        else:
            st.markdown(message["content"])
        if "timestamp" in message:
            st.caption(f"🕒 {message['timestamp']}")

# Manual HF Analysis Section
if st.session_state.messages and len(st.session_state.messages) > 0:
    st.divider()
    
    # HF Expert Section
    with st.expander("🤖 HF Expert Analysis", expanded=False):
        st.subheader("Deep Conversation Analysis")
        
        col1, col2 = st.columns([3, 1])
        with col1:
            st.markdown("""
                **HF Expert Features:**
                - Analyzes entire conversation history
                - Performs web research when needed
                - Provides deep insights and recommendations
                - Acts as expert consultant in your conversation
            """)
            
            # Show conversation preview
            st.markdown("**Conversation Preview for HF Expert:**")
            st.markdown("---")
            for i, msg in enumerate(st.session_state.messages[-5:]):  # Last 5 messages
                role = "👤 You" if msg["role"] == "user" else "🤖 Assistant"
                st.markdown(f"**{role}:** {msg['content'][:100]}{'...' if len(msg['content']) > 100 else ''}")
            st.markdown("---")
        
        with col2:
            if st.button("🧠 Activate HF Expert",
                         key="activate_hf_expert",
                         help="Send conversation to HF endpoint for deep analysis",
                         use_container_width=True,
                         disabled=st.session_state.is_processing):
                st.session_state.hf_expert_requested = True
        
        # Show HF expert analysis when requested
        if st.session_state.get("hf_expert_requested", False):
            with st.spinner("🧠 HF Expert analyzing conversation..."):
                try:
                    # Get conversation history
                    user_session = session_manager.get_session("default_user")
                    conversation_history = user_session.get("conversation", [])
                    
                    # Show what HF expert sees
                    with st.expander("📋 HF Expert Input", expanded=False):
                        st.markdown("**Conversation History Sent to HF Expert:**")
                        for i, msg in enumerate(conversation_history[-10:]):  # Last 10 messages
                            st.markdown(f"**{msg['role'].capitalize()}:** {msg['content'][:100]}{'...' if len(msg['content']) > 100 else ''}")
                    
                    # Request HF analysis
                    hf_analysis = coordinator.manual_hf_analysis(
                        "default_user",
                        conversation_history
                    )
                    
                    if hf_analysis:
                        # Display HF expert response with clear indication
                        with st.chat_message("assistant"):
                            st.markdown("### 🤖 HF Expert Analysis")
                            st.markdown(hf_analysis)
                        
                        # Add research/web search decisions
                        research_needs = coordinator.determine_web_search_needs(conversation_history)
                        if research_needs["needs_search"]:
                            st.info(f"🔍 **Research Needed:** {research_needs['reasoning']}")
                            if st.button("🔎 Perform Web Research", key="web_research_button"):
                                # Perform web search
                                with st.spinner("🔎 Searching for current information..."):
                                    # Add web search logic here
                                    st.success("✅ Web research completed!")
                        
                        # Add to message history with HF expert tag
                        st.session_state.messages.append({
                            "role": "assistant",
                            "content": hf_analysis,
                            "timestamp": datetime.now().strftime("%H:%M:%S"),
                            "source": "hf_expert",
                            "research_needs": research_needs
                        })
                        
                        st.session_state.hf_expert_requested = False
                
                except Exception as e:
                    user_msg = translate_error(e)
                    st.error(f"❌ HF Expert analysis failed: {user_msg}")
                    st.session_state.hf_expert_requested = False

# 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:
    st.session_state.is_processing = True
    
    # Display user message
    with st.chat_message("user"):
        st.markdown(user_input)
    
    st.session_state.messages.append({
        "role": "user",
        "content": user_input,
        "timestamp": datetime.now().strftime("%H:%M:%S")
    })
    
    # Process AI response
    with st.chat_message("assistant"):
        response_placeholder = st.empty()
        status_placeholder = st.empty()
        
        try:
            # Get conversation history
            user_session = session_manager.get_session("default_user")
            conversation = user_session.get("conversation", [])
            conversation_history = conversation[-5:]  # Last 5 messages
            conversation_history.append({"role": "user", "content": user_input})
            
            # Try Ollama with proper error handling
            status_placeholder.info(PROCESSING_STAGES["ollama"])
            ai_response = None
            
            try:
                ai_response = send_to_ollama(
                    user_input,
                    conversation_history,
                    st.session_state.ngrok_url_temp,
                    st.session_state.selected_model
                )
                
                if ai_response:
                    response_placeholder.markdown(ai_response)
                    status_placeholder.success("✅ Response received!")
                else:
                    status_placeholder.warning("⚠️ Empty response from Ollama")
            
            except Exception as ollama_error:
                user_msg = translate_error(ollama_error)
                status_placeholder.error(f"⚠️ {user_msg}")
            
            # Fallback to HF if available
            if config.hf_token and not ai_response:
                status_placeholder.info(PROCESSING_STAGES["hf_init"])
                try:
                    ai_response = send_to_hf(user_input, conversation_history)
                    if ai_response:
                        response_placeholder.markdown(ai_response)
                        status_placeholder.success("✅ HF response received!")
                    else:
                        status_placeholder.error("❌ No response from HF")
                except Exception as hf_error:
                    user_msg = translate_error(hf_error)
                    status_placeholder.error(f"⚠️ {user_msg}")
            
            # Save response if successful
            if ai_response:
                # Update conversation history
                conversation.append({"role": "user", "content": user_input})
                conversation.append({"role": "assistant", "content": ai_response})
                user_session["conversation"] = conversation
                session_manager.update_session("default_user", user_session)
                
                # Add to message history
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": ai_response,
                    "timestamp": datetime.now().strftime("%H:%M:%S")
                })
            else:
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": "Sorry, I couldn't process your request. Please try again.",
                    "timestamp": datetime.now().strftime("%H:%M:%S")
                })
        
        except Exception as e:
            user_msg = translate_error(e)
            response_placeholder.error(f"⚠️ {user_msg}")
            st.session_state.messages.append({
                "role": "assistant",
                "content": f"⚠️ {user_msg}",
                "timestamp": datetime.now().strftime("%H:%M:%S")
            })
        finally:
            st.session_state.is_processing = False
            time.sleep(0.5)  # Brief pause
            st.experimental_rerun()