src/session_manager.py

"""
Session management for the AI Database Assistant.
"""
import streamlit as st
from typing import List, Dict, Any


class SessionManager:
    """Manages session state and message handling."""
    
    @staticmethod
    def initialize_session() -> None:
        """Initialize session state variables."""
        if "messages" not in st.session_state:
            st.session_state["messages"] = []
    
    @staticmethod
    def get_messages() -> List[Dict[str, Any]]:
        """Get all messages from session state."""
        return st.session_state.get("messages", [])
    
    @staticmethod
    def add_user_message(content: str) -> None:
        """Add a user message to the session."""
        st.session_state["messages"].append({
            "role": "user",
            "content": content
        })
    
    @staticmethod
    def add_thinking_message() -> None:
        """Add a thinking placeholder message."""
        st.session_state["messages"].append({
            "role": "assistant",
            "content": "🤔 Thinking...",
            "is_placeholder": True
        })
    
    @staticmethod
    def replace_last_message(content: str, data: List[Dict] = None, 
                           chart: str = None, is_error: bool = False) -> None:
        """Replace the last message with final response."""
        st.session_state["messages"][-1] = {
            "role": "assistant",
            "content": content,
            "data": data or [],
            "chart": chart,
            "is_error": is_error,
        }
    
    @staticmethod
    def is_ai_thinking() -> bool:
        """Check if AI is currently thinking (has placeholder message)."""
        messages = SessionManager.get_messages()
        if not messages:
            return False
        
        last_message = messages[-1]
        return (
            last_message["role"] == "assistant" and 
            last_message.get("is_placeholder", False)
        )
    
    @staticmethod
    def get_last_user_message() -> str:
        """Get the content of the last user message."""
        messages = SessionManager.get_messages()
        if len(messages) >= 2:
            return messages[-2]["content"]
        return ""
    
    @staticmethod
    def has_pending_response() -> bool:
        """Check if there's a pending response to process."""
        messages = SessionManager.get_messages()
        return (
            messages and
            messages[-1].get("is_placeholder") and
            len(messages) >= 2 and
            messages[-2]["role"] == "user"
        )