streamlit_app.py

import streamlit as st
import requests
import json
import uuid
from datetime import datetime
import os

BACKEND_URL = os.getenv("BACKEND_URL", "http://localhost:8001")
BACKEND_DB = "/app/chat_history.db"  # Same DB used by backend

# Page config
st.set_page_config(
    page_title="AI Agent Chat",
    page_icon="🤖",
    layout="wide"
)

# Custom CSS
st.markdown("""
    <style>
    .main {
        padding: 0rem 1rem;
    }
    .stChatMessage {
        padding: 1rem;
        border-radius: 0.5rem;
    }
    </style>
""", unsafe_allow_html=True)


# REMOVE get_all_sessions_from_db() - use API instead
def get_all_sessions_from_api():
    """Get all sessions from backend API"""
    try:
        # Add this endpoint to backend if not exists
        response = requests.get(f"{BACKEND_URL}/sessions", timeout=5)
        if response.status_code == 200:
            return response.json().get("sessions", [])
        return []
    except Exception as e:
        st.error(f"Failed to fetch sessions: {e}")
        return []


# REMOVE get_session_preview() - use API instead
def get_session_preview_from_api(session_id: str):
    """Get session preview from backend API"""
    try:
        response = requests.get(
            f"{BACKEND_URL}/history/{session_id}",
            params={"limit": 1},
            timeout=5
        )
        if response.status_code == 200:
            data = response.json()
            messages = data.get("messages", [])
            if messages and messages[0]["role"] == "user":
                preview = messages[0]["content"]
                return preview[:50] + "..." if len(preview) > 50 else preview
        return "Empty chat"
    except Exception:
        return "Unknown"


def check_backend_health():
    """Check if backend is running"""
    try:
        response = requests.get(f"{BACKEND_URL}/health", timeout=5)
        return response.status_code == 200
    except Exception:
        return False


def get_chat_history(session_id: str):
    """Fetch chat history from backend"""
    try:
        response = requests.get(
            f"{BACKEND_URL}/history/{session_id}",
            timeout=5
        )
        if response.status_code == 200:
            data = response.json()
            return data.get("messages", [])
        return []
    except Exception as e:
        st.error(f"Failed to fetch history: {e}")
        return []


def clear_backend_history(session_id: str):
    """Clear chat history from backend"""
    try:
        response = requests.delete(
            f"{BACKEND_URL}/history/{session_id}",
            timeout=5
        )
        return response.status_code == 200
    except Exception:
        return False


def load_session(session_id: str):
    """Load a specific session"""
    st.session_state.session_id = session_id
    st.session_state.messages = []
    
    # Load history from backend
    history = get_chat_history(session_id)
    if history:
        for msg in history:
            st.session_state.messages.append({
                "role": msg["role"],
                "content": msg["content"]
            })
    
    st.rerun()


def create_new_session():
    """Create a new session"""
    new_session_id = str(uuid.uuid4())
    st.session_state.session_id = new_session_id
    st.session_state.messages = []
    st.rerun()


def send_message_streaming(message: str):
    """Send message and stream response TOKEN BY TOKEN"""
    try:
        url = f"{BACKEND_URL}/chat/stream"
        payload = {
            "message": message,
            "session_id": st.session_state.session_id
        }
        
        response = requests.post(
            url,
            json=payload,
            stream=True,
            timeout=60
        )
        
        if response.status_code != 200:
            return f"Error: {response.status_code}"
        
        # Placeholders for streaming
        tool_placeholder = st.empty()
        message_placeholder = st.empty()
        
        full_response = ""
        tool_calls = []
        tool_results = []
        current_tools_shown = False
        
        # Process stream line by line
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    try:
                        data = json.loads(line[6:])
                        
                        if data.get("type") == "tool_call":
                            tools = data.get("tools", [])
                            tool_calls.extend(tools)
                            tool_placeholder.markdown(f"🔧 **Using tools:** {', '.join(tools)}")
                            current_tools_shown = True
                        
                        elif data.get("type") == "tool_start":
                            tool_name = data.get("tool", "unknown")
                            tool_placeholder.markdown(f"⚙️ **Executing:** {tool_name}...")
                        
                        elif data.get("type") == "tool_result":
                            tool_name = data.get("tool", "unknown")
                            preview = data.get("preview", "")
                            tool_results.append(f"{tool_name}: {preview[:100]}")
                            tool_placeholder.markdown(f"✅ **Completed:** {tool_name}")
                        
                        elif data.get("type") == "token":
                            # TRUE TOKEN-BY-TOKEN STREAMING
                            full_response = data.get("content", "")
                            message_placeholder.markdown(full_response + "▌")  # Show cursor
                        
                        elif data.get("type") == "content":
                            # Fallback for complete content
                            full_response = data.get("content", "")
                            message_placeholder.markdown(full_response + "▌")
                        
                        elif data.get("type") == "done":
                            # Remove cursor and clear tool placeholder
                            if full_response:
                                message_placeholder.markdown(full_response)
                            else:
                                message_placeholder.warning("No response generated. The tool may have encountered an issue.")
                            
                            if current_tools_shown:
                                tool_placeholder.empty()
                            break
                        
                        elif data.get("type") == "error":
                            st.error(f"Error: {data.get('error')}")
                            return None
                    
                    except json.JSONDecodeError:
                        continue
        
        return full_response
    
    except Exception as e:
        st.error(f"Failed to send message: {e}")
        return None


# Initialize session state
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())

if "messages" not in st.session_state:
    st.session_state.messages = []


# Sidebar
with st.sidebar:
    st.title("🤖 AI Agent")
    
    # Check backend status
    is_healthy = check_backend_health()
    
    if not is_healthy:
        st.error("❌ Backend not available")
        st.info(f"Backend URL: {BACKEND_URL}")
        st.warning("Check if backend Space is running")
    else:
        st.success("✅ Connected to backend")
    
    st.divider()
    
    # New Chat Button
    if st.button("➕ New Chat", use_container_width=True, type="primary"):
        create_new_session()
    
    st.divider()
    
    # Session History - now from API
    st.subheader("💬 Chat History")
    
    sessions = get_all_sessions_from_api()
    
    if sessions:
        for session in sessions:
            session_id = session.get("session_id")
            preview = session.get("preview", "Unknown")
            is_current = session_id == st.session_state.session_id
            
            col1, col2 = st.columns([5, 1])
            
            with col1:
                button_label = f"{'📍' if is_current else '💬'} {preview}"
                if st.button(
                    button_label,
                    key=f"load_{session_id}",
                    use_container_width=True,
                    disabled=is_current,
                ):
                    load_session(session_id)
            
            with col2:
                if len(sessions) > 1 or not is_current:
                    if st.button("🗑️", key=f"delete_{session_id}", help="Delete"):
                        if clear_backend_history(session_id):
                            if is_current:
                                create_new_session()
                            else:
                                st.rerun()
    else:
        st.info("No chat history yet")
    
    st.divider()
    
    # Current Session Info
    with st.expander("ℹ️ Current Session"):
        st.text(f"ID: {st.session_state.session_id[:16]}...")
        st.text(f"Messages: {len(st.session_state.messages)}")
        st.text(f"Backend: {BACKEND_URL}")
        
        if st.button("🗑️ Clear Current Chat", use_container_width=True):
            if clear_backend_history(st.session_state.session_id):
                st.session_state.messages = []
                st.success("Chat cleared!")
                st.rerun()
    
    st.divider()
    
    # Info
    with st.expander("🔧 Features"):
        st.markdown("""
        - 💬 **Streaming responses**
        - 🧠 **Conversation memory**
        - 🔍 **Web search**
        - 🌤️ **Weather info**
        - 💱 **Currency conversion**
        - 📚 **Wikipedia lookup**
        - ⏰ **World time**
        """)


# Main chat interface
st.title("💬 Chat with AI Agent")

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

# Chat input
if prompt := st.chat_input("Type your message here...", disabled=not is_healthy):
    # Add user message to chat
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Get streaming response
    with st.chat_message("assistant"):
        response = send_message_streaming(prompt)
        
        if response:
            st.session_state.messages.append({
                "role": "assistant",
                "content": response
            })

# Footer
st.divider()
st.caption(f"🔗 Backend: {BACKEND_URL}")