app.py

import streamlit as st
import requests
import os
import json
from datetime import datetime
import streamlit_cookies_manager as scm
from config import get_api_config, validate_url, is_configured, get_config_status, get_app_title, get_secret
import uuid

# Page config
st.set_page_config(
    page_title="Ejento AI Chat",
    page_icon="💬",
    layout="centered"
)

# Get API config for cookie manager (before session state)
_temp_config = get_api_config()

# Initialize cookie manager
# Use ACCESS_TOKEN as cookie encryption password for consistency
cookies = scm.EncryptedCookieManager(
    prefix="ejento_chat_",
    password=_temp_config['access_token'] or "ejento-chat-default-key"
)

if not cookies.ready():
    st.stop()

# Custom CSS to make thread list buttons look like plain text
st.markdown("""
<style>
/* Move main app title up */
.main .block-container {
    padding-top: 0.5rem !important;
}

h1 {
    margin-top: 0rem !important;
    padding-top: 0rem !important;
}
/* Remove all button styling for sidebar buttons */
.stSidebar button[kind="secondary"],
.stSidebar button[data-testid="baseButton-secondary"] {
    background: none !important;
    background-color: transparent !important;
    border: none !important;
    box-shadow: none !important;
    padding: 0.5rem 0rem !important;
    margin: 0rem !important;
    text-align: left !important;
    color: inherit !important;
    font-size: 16px !important;
    justify-content: flex-start !important;
    width: 100% !important;
    display: block !important;
}

.stSidebar button[kind="secondary"]:hover,
.stSidebar button[data-testid="baseButton-secondary"]:hover {
    background-color: rgba(151, 166, 195, 0.15) !important;
    border: none !important;
    text-decoration: underline !important;
}

.stSidebar button[kind="secondary"]:focus,
.stSidebar button[data-testid="baseButton-secondary"]:focus {
    background: none !important;
    border: none !important;
    box-shadow: none !important;
}

/* Make markdown text in sidebar bigger and align left */
.stSidebar .stMarkdown p {
    font-size: 16px !important;
    margin: 0rem !important;
    padding: 0.5rem 0rem !important;
}

/* Target button text content directly */
.stSidebar button[kind="secondary"] p,
.stSidebar button[kind="secondary"] div,
.stSidebar button[kind="secondary"] span,
.stSidebar button[data-testid="baseButton-secondary"] p,
.stSidebar button[data-testid="baseButton-secondary"] div,
.stSidebar button[data-testid="baseButton-secondary"] span {
    font-size: 16px !important;
}

/* Also target all text in sidebar */
.stSidebar * {
    font-size: 16px !important;
}

/* Reduce spacing around divider */
.stSidebar hr {
    margin-top: 0.5rem !important;
    margin-bottom: 0.5rem !important;
}

/* Chat message content font size */
[data-testid="stChatMessageContent"] p {
    font-size: 16px !important;
}

/* Chat input field font size */
[data-testid="stChatInput"] textarea {
    font-size: 16px !important;
}

[data-testid="stChatInput"] textarea::placeholder {
    font-size: 16px !important;
}

[data-testid="stChatInput"] * {
    font-size: 16px !important;
}

.stChatInput textarea {
    font-size: 16px !important;
}

textarea {
    font-size: 16px !important;
}
</style>
""", unsafe_allow_html=True)

# Initialize session state
if 'messages' not in st.session_state:
    st.session_state.messages = []

if 'current_thread_id' not in st.session_state:
    st.session_state.current_thread_id = None

if 'thread_title' not in st.session_state:
    st.session_state.thread_title = None

if 'thread_loaded' not in st.session_state:
    st.session_state.thread_loaded = False

if 'api_config' not in st.session_state:
    st.session_state.api_config = get_api_config()
    st.session_state.api_config['base_url'] = validate_url(st.session_state.api_config['base_url'])

# Message tracking to prevent duplicates - using message IDs only
if 'last_sent_message_id' not in st.session_state:
    st.session_state.last_sent_message_id = None

if 'is_sending' not in st.session_state:
    st.session_state.is_sending = False

if 'pending_message' not in st.session_state:
    st.session_state.pending_message = None

if 'current_message_id' not in st.session_state:
    st.session_state.current_message_id = None

# API Functions
def create_thread(title=None):
    """Create a new chat thread with optional title"""
    base_url = validate_url(st.session_state.api_config['base_url'])
    url = f"{base_url}/api/v2/agents/{st.session_state.api_config['agent_id']}/chat-threads"
    
    headers = {
        'Ocp-Apim-Subscription-Key': st.session_state.api_config['api_key'],
        'Authorization': f"Bearer {st.session_state.api_config['access_token']}",
        'Content-Type': 'application/json'
    }
    
    # Use provided title or default to timestamp
    if not title:
        title = f"New Chat"
    
    payload = {
        "title": title
    }
    
    try:
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        data = response.json()
        if data.get('success'):
            thread_id = data['data']['id']
            st.session_state.thread_title = title
            st.success(f"✅ Thread created: {title}")
            return thread_id, title
        else:
            st.error(f"❌ Failed to create thread: {data.get('message', 'Unknown error')}")
    except requests.exceptions.RequestException as e:
        st.error(f"❌ Error creating thread: {str(e)}")
        if hasattr(e, 'response') and e.response is not None:
            try:
                error_data = e.response.json()
                st.error(f"API Error: {error_data}")
            except:
                st.error(f"Response: {e.response.text}")
    except Exception as e:
        st.error(f"❌ Unexpected error: {str(e)}")
    return None, None

def update_thread_title(thread_id, new_title):
    """Update the title of an existing chat thread"""
    base_url = validate_url(st.session_state.api_config['base_url'])
    url = f"{base_url}/api/v2/chat-threads/{thread_id}"
    
    headers = {
        'Ocp-Apim-Subscription-Key': st.session_state.api_config['api_key'],
        'Authorization': f"Bearer {st.session_state.api_config['access_token']}",
        'Content-Type': 'application/json'
    }
    
    payload = {
        "title": new_title
    }
    
    try:
        response = requests.put(url, json=payload, headers=headers)
        response.raise_for_status()
        data = response.json()
        
        if data.get('success'):
            return True
        else:
            st.warning(f"Failed to update thread title: {data.get('message', 'Unknown error')}")
            return False
    except Exception as e:
        st.warning(f"Could not update thread title: {str(e)}")
        return False

def get_thread_details(thread_id):
    """Fetch details for a specific thread from API"""
    try:
        base_url = validate_url(st.session_state.api_config['base_url'])
        url = f"{base_url}/api/v2/chat-threads/{thread_id}"
        
        headers = {
            'Ocp-Apim-Subscription-Key': st.session_state.api_config['api_key'],
            'Authorization': f"Bearer {st.session_state.api_config['access_token']}"
        }
        
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        
        data = response.json()
        if data and 'data' in data:
            thread_data = data['data']
            return {
                'id': thread_data.get('id'),
                'title': thread_data.get('title') or thread_data.get('name') or 'New Chat'
            }
        return None
    except Exception as e:
        st.warning(f"Could not fetch thread details: {str(e)}")
        return None

def get_all_threads():
    """Fetch all chat threads from API"""
    base_url = validate_url(st.session_state.api_config['base_url'])
    url = f"{base_url}/api/v2/agents/{st.session_state.api_config['agent_id']}/chat-threads"
    
    headers = {
        'Ocp-Apim-Subscription-Key': st.session_state.api_config['api_key'],
        'Authorization': f"Bearer {st.session_state.api_config['access_token']}"
    }
    
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        data = response.json()
        
        if isinstance(data, dict) and data.get('success'):
            # Response format: {"success": true, "data": {"chat_threads": [...]}}
            data_obj = data.get('data', {})
            if isinstance(data_obj, dict):
                threads = data_obj.get('chat_threads', [])
            else:
                threads = []
        elif isinstance(data, list):
            threads = data
        else:
            return []
        
        # Convert to our format and sort by created_on
        thread_list = []
        for thread in threads:
            if isinstance(thread, dict):
                thread_list.append({
                    'id': thread.get('id'),
                    'title': thread.get('title', 'Untitled'),
                    'created_on': thread.get('created_on', datetime.now().isoformat())
                })
        
        # Sort by created_on (most recent first)
        thread_list.sort(key=lambda x: x.get('created_on', ''), reverse=True)
        
        return thread_list
        
    except Exception as e:
        st.warning(f"Could not load threads: {str(e)}")
        return []

def get_chat_logs(thread_id):
    """Retrieve chat logs for a specific thread"""
    base_url = validate_url(st.session_state.api_config['base_url'])
    url = f"{base_url}/api/v2/chat-threads/{thread_id}/chat-logs"
    
    headers = {
        'Ocp-Apim-Subscription-Key': st.session_state.api_config['api_key'],
        'Authorization': f"Bearer {st.session_state.api_config['access_token']}"
    }
    
    try:
        response = requests.get(url, headers=headers)
        response.raise_for_status()
        data = response.json()
        

        # Handle API response structure
        if isinstance(data, dict) and data.get('success'):
            # Response format: {"success": true, "data": {"chat_logs": [...]}}
            data_obj = data.get('data', {})
            if isinstance(data_obj, dict):
                chat_logs = data_obj.get('chat_logs', [])
            else:
                chat_logs = []
        elif isinstance(data, list):
            # API returned list directly
            chat_logs = data
        else:
            return []
        
        messages = []
        
        # Ensure chat_logs is a list
        if not isinstance(chat_logs, list):
            return []
        
        # Sort chat_logs by created_on (oldest first) or by id
        chat_logs_sorted = sorted(chat_logs, key=lambda x: (x.get('created_on', ''), x.get('id', 0)))
        
        # Deduplicate log entries by ID (keep first occurrence of each unique log ID)
        seen_log_ids = set()
        deduplicated_logs = []
        for log in chat_logs_sorted:
            if not isinstance(log, dict):
                continue
            log_id = log.get('id')
            if log_id and log_id not in seen_log_ids:
                seen_log_ids.add(log_id)
                deduplicated_logs.append(log)
            elif not log_id:
                # If no ID, include it anyway (shouldn't happen but be safe)
                deduplicated_logs.append(log)
        
        # Now process deduplicated logs into messages
        for log in deduplicated_logs:
            # Add user message
            if log.get('question'):
                messages.append({
                    "role": "user",
                    "content": log['question'],
                    "timestamp": log.get('created_on', datetime.now().isoformat())
                })
            
            # Add assistant message - field is 'response'
            if log.get('response'):
                response_data = log['response']
                
                # Handle both string and dict formats
                if isinstance(response_data, dict):
                    # Response is a dict like {'answer': '...'}
                    content = response_data.get('answer', str(response_data))
                elif isinstance(response_data, str):
                    # Response is already a string
                    content = response_data
                else:
                    # Unknown format, convert to string
                    content = str(response_data)
                
                messages.append({
                    "role": "assistant",
                    "content": content,
                    "timestamp": log.get('created_on', datetime.now().isoformat())
                })
        
        # Final check: ensure proper user-bot alternation
        # Remove consecutive messages of the same role
        final_messages = []
        for msg in messages:
            # Add if list is empty or role is different from last message
            if not final_messages or final_messages[-1]['role'] != msg['role']:
                final_messages.append(msg)
            # If same role, keep the first one (skip this duplicate)
        
        return final_messages
        
    except requests.exceptions.RequestException as e:
        st.warning(f"Could not load chat history: API error - {str(e)}")
        return []
    except json.JSONDecodeError as e:
        st.warning(f"Could not load chat history: Invalid JSON response")
        return []
    except Exception as e:
        st.warning(f"Could not load chat history: {str(e)}")
        return []

def send_message_stream(message, thread_id, message_id, placeholder):
    """Send message and stream the response - with duplicate prevention using message ID
    Returns: (answer, thread_id, thread_title) tuple
    """
    
    # Check if this message ID has already been sent
    if st.session_state.last_sent_message_id == message_id:
        # This message was already sent, skip
        return ("ALREADY_SENT", None, None)
    
    try:
        base_url = validate_url(st.session_state.api_config['base_url'])
        url = f"{base_url}/response-service/api/v2/agents/{st.session_state.api_config['agent_id']}/responses/stream"
        
        headers = {
            'Ocp-Apim-Subscription-Key': st.session_state.api_config['api_key'],
            'Authorization': f"Bearer {st.session_state.api_config['access_token']}",
            'Content-Type': 'application/json'
        }
        
        # Build history from previous messages (exclude current unpaired message)
        history = []
        messages = st.session_state.messages
        for i in range(0, len(messages) - 1, 2):
            if i + 1 < len(messages):
                history.append({
                    "user": messages[i]["content"],
                    "bot": messages[i + 1]["content"]
                })
        
        payload = {
            "user_query": message,
            "query_source": st.session_state.api_config['query_source'],
            "chat_thread_id": thread_id
        }
        
        if history:
            payload["history"] = history
        
        # Make streaming API call
        response = requests.post(url, json=payload, headers=headers, stream=True)
        response.raise_for_status()
        
        accumulated_text = ""
        current_event = None
        
        # Process SSE stream
        for line in response.iter_lines():
            if line:
                line_str = line.decode('utf-8')
                
                if line_str.startswith('event:'):
                    current_event = line_str.split(':', 1)[1].strip()
                    
                elif line_str.startswith('data:'):
                    data_str = line_str.split(':', 1)[1].strip()
                    
                    try:
                        data = json.loads(data_str)
                        
                        if current_event == 'token':
                            # Stream token in real-time
                            delta = data.get('delta', '')
                            accumulated_text += delta
                            # Update placeholder with streaming cursor
                            placeholder.markdown(accumulated_text + " █")
                            
                        elif current_event == 'end':
                            # Stream complete - remove cursor
                            placeholder.markdown(accumulated_text)
                            st.session_state.last_sent_message_id = message_id
                            
                            # Extract thread_id and thread_title from response if available
                            response_thread_id = None
                            response_thread_title = None
                            
                            if 'output' in data:
                                output = data['output']
                                response_thread_id = output.get('thread_id')
                                response_thread_title = output.get('thread_title') or output.get('title')
                            elif 'thread_id' in data:
                                response_thread_id = data.get('thread_id')
                                response_thread_title = data.get('thread_title') or data.get('title')
                            
                            return (accumulated_text, response_thread_id, response_thread_title)
                            
                        elif current_event == 'error':
                            error_msg = data.get('message', 'Unknown error')
                            placeholder.error(f"API Error: {error_msg}")
                            return (None, None, None)
                            
                    except json.JSONDecodeError:
                        pass
        
        # Fallback if stream ends without 'end' event
        if accumulated_text:
            placeholder.markdown(accumulated_text)
            st.session_state.last_sent_message_id = message_id
            return (accumulated_text, None, None)
        else:
            placeholder.warning("No response received")
            return (None, None, None)
        
    except requests.exceptions.RequestException as e:
        error_msg = f"API Error: {str(e)}"
        if hasattr(e, 'response') and e.response is not None:
            try:
                error_data = e.response.json()
                error_msg += f" - {error_data}"
            except:
                error_msg += f" - {e.response.text[:200]}"
        placeholder.error(error_msg)
        return (None, None, None)
    except Exception as e:
        st.error(f"Unexpected error: {str(e)}")
        import traceback
        st.error(traceback.format_exc())
        return (None, None, None)

# Cookie-based thread persistence
def save_thread_to_cookie(thread_id, title):
    """Save thread ID and title to browser cookie"""
    try:
        cookies['thread_id'] = str(thread_id)
        cookies['thread_title'] = title
        cookies.save()
    except Exception as e:
        # Silently ignore all cookie save errors
        pass

def load_thread_from_cookie():
    """Load thread ID and title from browser cookie"""
    try:
        thread_id = cookies.get('thread_id')
        thread_title = cookies.get('thread_title')
        
        if thread_id:
            return {
                'thread_id': int(thread_id),
                'title': thread_title
            }
    except Exception as e:
        st.warning(f"Could not load thread from cookie: {str(e)}")
    return None

# Thread list management in cookies
def get_thread_list():
    """Get list of threads from cookies"""
    try:
        thread_list_json = cookies.get('thread_list')
        if thread_list_json:
            thread_list = json.loads(thread_list_json)
            
            # Deduplicate on retrieval (safety check)
            seen_ids = set()
            deduplicated = []
            for thread in thread_list:
                tid = int(thread['id'])
                if tid not in seen_ids:
                    seen_ids.add(tid)
                    deduplicated.append(thread)
            
            return deduplicated
        return []
    except Exception as e:
        return []

def save_thread_list(thread_list):
    """Save thread list to cookies"""
    try:
        cookies['thread_list'] = json.dumps(thread_list)
        cookies.save()
    except Exception as e:
        pass

def add_thread_to_list(thread_id, title):
    """Add a new thread to the cookie-based thread list"""
    thread_list = get_thread_list()
    
    # Check if thread already exists (by ID)
    existing_ids = {int(t['id']) for t in thread_list}
    
    if int(thread_id) not in existing_ids:
        thread_list.insert(0, {
            'id': int(thread_id),
            'title': title,
            'created_on': datetime.now().isoformat()
        })
        save_thread_list(thread_list)

def load_thread(thread_id, title):
    """Load a specific thread"""
    st.session_state.current_thread_id = thread_id
    st.session_state.thread_title = title
    st.session_state.thread_loaded = False
    
    # Load messages from API
    messages = get_chat_logs(thread_id)
    st.session_state.messages = messages
    
    # Save to cookie
    save_thread_to_cookie(thread_id, title)
    
    # Clear any pending messages and reset message ID tracking
    st.session_state.pending_message = None
    st.session_state.current_message_id = None
    st.session_state.last_sent_message_id = None
    st.session_state.is_sending = False
    
    st.session_state.thread_loaded = True

# Sidebar - Thread Management
with st.sidebar:
    # New Chat button
    if st.button("➕ New Chat", use_container_width=True, type="secondary"):
        # Clear session state
        st.session_state.current_thread_id = None
        st.session_state.thread_title = None
        st.session_state.messages = []
        st.session_state.thread_loaded = False
        st.session_state.pending_message = None
        st.session_state.current_message_id = None
        st.session_state.last_sent_message_id = None
        st.session_state.is_sending = False
        
        # Clear cookie so it doesn't reload on rerun
        try:
            cookies['thread_id'] = ''
            cookies['thread_title'] = ''
            cookies.save()
        except:
            pass
        
        st.rerun()
    
    st.divider()
    
    # Display thread list from cookies
    thread_list = get_thread_list()
    
    if thread_list:
        st.markdown("**Recent Threads:**")
        for thread in thread_list:
            thread_id = int(thread['id'])
            thread_title = thread.get('title', 'Untitled')
            
            # Highlight selected thread
            if st.session_state.current_thread_id == thread_id:
                st.markdown(f"<div style='background-color: rgba(128, 128, 128, 0.2); padding: 0.5rem; border-radius: 0.25rem; margin: 0.25rem 0; font-size: 0.875rem;'>{thread_title}</div>", unsafe_allow_html=True)
            else:
                if st.button(thread_title, key=f"thread_{thread_id}", use_container_width=True, type="secondary"):
                    load_thread(thread_id, thread_title)
                    st.rerun()
    else:
        st.markdown("*No threads yet. Start a new chat!*")

# Main chat area
# Ensure thread_loaded flag is properly set
if not st.session_state.thread_loaded:
    st.session_state.thread_loaded = True

# Always start with new chat on page reload
# Don't load thread from cookie - user can manually load threads from sidebar
if st.session_state.current_thread_id is None:
    # Leave thread_id as None
    # User will start a new chat
    # API will create thread on first message and return thread_id
    pass

# Custom CSS - only for sidebar thread buttons
st.markdown("""
<style>
    /* Make sidebar buttons smaller */
    [data-testid="stSidebar"] .stButton button {
        font-size: 0.875rem !important;
    }
</style>
""", unsafe_allow_html=True)

# Title
st.title(get_app_title())

# Display thread info or greeting
if st.session_state.current_thread_id:
    st.caption(f"📝 {st.session_state.thread_title}")
    st.divider()
else:
    # Simple new chat indicator
    st.subheader("👋 New Chat")
    st.divider()

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

# Process pending message if exists
if st.session_state.pending_message and st.session_state.current_message_id and not st.session_state.is_sending:
    # Allow sending even if thread_id is None (API will create thread)
    if True:
        # Set sending flag to prevent duplicate processing during reruns
        st.session_state.is_sending = True
        
        # Show loading indicator while processing
        with st.spinner("Thinking... ⏳"):
            try:
                # Create a temporary placeholder just for status
                status_placeholder = st.empty()
                status_placeholder.info(f"Processing your message... (Thread ID: {st.session_state.current_thread_id or 'None - will be created'})")
                
                result = send_message_stream(
                    st.session_state.pending_message, 
                    st.session_state.current_thread_id,
                    st.session_state.current_message_id,
                    status_placeholder
                )
                
                # Unpack result
                answer, response_thread_id, response_thread_title = result
                
                # If we got a thread_id from the response and don't have one yet, use it
                if response_thread_id and not st.session_state.current_thread_id:
                    st.session_state.current_thread_id = response_thread_id
                    
                    # Fetch thread details from API to get proper title
                    thread_details = get_thread_details(response_thread_id)
                    if thread_details:
                        thread_title = thread_details['title']
                    else:
                        # Fallback to response title or "New Chat"
                        thread_title = response_thread_title or "New Chat"
                    
                    st.session_state.thread_title = thread_title
                    save_thread_to_cookie(response_thread_id, thread_title)
                    add_thread_to_list(response_thread_id, thread_title)
                elif response_thread_id:
                    pass
                
                # Store result for processing after spinner closes
                st.session_state.last_answer = answer
                
            except Exception as e:
                st.session_state.last_answer = None
                st.session_state.last_error = str(e)
            
            finally:
                # Always clear flags and pending message
                st.session_state.is_sending = False
                st.session_state.pending_message = None
                st.session_state.current_message_id = None
        
        # Process answer AFTER spinner closes
        if hasattr(st.session_state, 'last_answer'):
            answer = st.session_state.last_answer
            if answer and answer != "ALREADY_SENT":
                # Add to messages
                new_message = {
                    "role": "assistant",
                    "content": answer,
                    "timestamp": datetime.now().isoformat()
                }
                st.session_state.messages.append(new_message)
            elif not answer or answer == "ALREADY_SENT":
                # Show error
                if answer == "ALREADY_SENT":
                    st.warning("Message already sent")
                else:
                    st.error("⚠️ No response received from API")
            
            # Clean up
            del st.session_state.last_answer
        
        if hasattr(st.session_state, 'last_error'):
            st.error(f"❌ Error: {st.session_state.last_error}")
            del st.session_state.last_error
        
        # Rerun to display the message from messages loop
        st.rerun()

# Chat input
if prompt := st.chat_input("Type your message..."):
    # Generate unique message ID
    message_id = str(uuid.uuid4())
    
    # Thread already exists with hardcoded name, no need to update
    
    # Add user message to chat
    st.session_state.messages.append({
        "role": "user",
        "content": prompt,
        "timestamp": datetime.now().isoformat(),
        "message_id": message_id
    })
    
    # Set pending message to be sent to API
    st.session_state.pending_message = prompt
    st.session_state.current_message_id = message_id
    
    st.rerun()