app.py

"""
Streamlit Frontend для RAG вопросно-ответной системы
Чат-интерфейс с поддержкой нескольких диалогов
"""
import streamlit as st
from datetime import datetime, timedelta
from typing import List, Dict, Optional
import uuid

from src import RAG
from src.db_utils.history_utils import (
    init_history_table,
    log_query,
    get_all_history,
    get_history_by_dialogue,
    search_history,
    get_history_stats,
    delete_history,
    get_recent_dialogues
)


# --- Инициализация RAG и БД ---
@st.cache_resource(show_spinner=False)
def get_rag():
    """Initialize RAG once and cache it"""
    return RAG(
        embed_model_name = "Qwen/Qwen3-Embedding-0.6B",
        embed_index_name = "recursive_Qwen3-Embedding-0.6B"
    )


@st.cache_resource(show_spinner=False)
def init_db():
    """Initialize database once and cache it"""
    try:
        init_history_table()
        return True
    except Exception as e:
        st.error(f"⚠️ Не удалось инициализировать таблицу истории: {e}")
        return False


# --- Session State Management ---
def init_session_state():
    """Initialize session state with caching"""
    if "current_dialogue_id" not in st.session_state:
        st.session_state.current_dialogue_id = None
    if "chat_list" not in st.session_state:
        st.session_state.chat_list = []
    if "current_chat_messages" not in st.session_state:
        st.session_state.current_chat_messages = []
    if "chat_list_loaded" not in st.session_state:
        st.session_state.chat_list_loaded = False


def generate_dialogue_id() -> str:
    """Generate unique dialogue ID"""
    return f"chat_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:6]}"


def get_chat_display_name(dialogue_id: str, first_query: str = None) -> str:
    """Get display name for chat - always from DB, no caching"""
    if first_query:
        # Use first 40 chars of first query as name
        name = first_query[:40] + "..." if len(first_query) > 40 else first_query
        return name
    
    return "Новый диалог"


# --- Chat Management Functions ---

def load_chats_list():
    """Load and cache chats list from DB"""
    try:
        st.session_state.chat_list = get_recent_dialogues(limit=50)
        st.session_state.chat_list_loaded = True
    except Exception as e:
        st.error(f"❌ Ошибка при загрузке чатов: {e}")
        st.session_state.chat_list = []


def create_new_chat():
    """Create a new chat"""
    new_id = generate_dialogue_id()
    st.session_state.current_dialogue_id = new_id
    st.session_state.current_chat_messages = []
    st.session_state.needs_rerun = True
    return new_id


def switch_to_chat(dialogue_id: str):
    """Switch to an existing chat and load its messages"""
    st.session_state.current_dialogue_id = dialogue_id
    load_current_chat_messages()
    st.session_state.needs_rerun = True


def load_current_chat_messages():
    """Load messages for current chat from DB and cache"""
    if not st.session_state.current_dialogue_id:
        st.session_state.current_chat_messages = []
        return
    
    try:
        st.session_state.current_chat_messages = get_history_by_dialogue(
            st.session_state.current_dialogue_id
        )
    except Exception as e:
        st.error(f"❌ Ошибка при загрузке сообщений: {e}")
        st.session_state.current_chat_messages = []


def get_current_chat_messages() -> List[Dict]:
    """Get cached messages for current chat"""
    return st.session_state.current_chat_messages


def send_message(query: str) -> Optional[Dict]:
    """Send a message in current chat and update cache"""
    try:
        if not st.session_state.current_dialogue_id:
            create_new_chat()
        
        # Get RAG and invoke with cached history
        rag = get_rag()
        
        # Use cached messages
        current_history = get_current_chat_messages()
        
        # Pass history to RAG (it will use last N messages internally for enrichment)
        result = rag.invoke(query, history=current_history)
        
        # Log to history DB
        query_id = log_query(
            query=query,
            answer=result.get("answer", ""),
            reason=result.get("reason", ""),
            dialogue_id=st.session_state.current_dialogue_id
        )
        
        result["query_id"] = query_id
        
        # Update only current messages, not all chats
        load_current_chat_messages()
        
        # Mark that we need to refresh chat list (but don't do it immediately)
        st.session_state.chat_list_loaded = False
        st.session_state.needs_rerun = True
        
        return result
    except Exception as e:
        st.error(f"❌ Ошибка при отправке сообщения: {e}")
        return None


def delete_chat(dialogue_id: str) -> bool:
    """Delete a chat from DB and update cache"""
    try:
        delete_history(dialogue_id=dialogue_id)
        
        # If deleted current chat, clear selection
        if st.session_state.current_dialogue_id == dialogue_id:
            st.session_state.current_dialogue_id = None
            st.session_state.current_chat_messages = []
        
        # Mark that we need to reload chat list
        st.session_state.chat_list_loaded = False
        st.session_state.needs_rerun = True
        
        return True
    except Exception as e:
        st.error(f"❌ Ошибка при удалении чата: {e}")
        return False




# --- Page: Chat Interface ---
def page_chat():
    """Main chat interface page"""
    
    # Custom CSS to fix chat input at the bottom + keyboard shortcuts
    st.markdown("""
        <style>
        /* Fix chat input at the bottom of main content area */
        section[data-testid="stSidebar"] ~ div .stChatInput {
            position: fixed;
            bottom: 0;
            background: white;
            padding: 1rem;
            z-index: 999;
            border-top: 1px solid #e6e6e6;
            margin-left: 0;
        }
        
        /* Add padding to main content to prevent overlap with fixed input */
        .main .block-container {
            padding-bottom: 100px;
        }
        
        /* Dark mode support */
        [data-testid="stAppViewContainer"][data-theme="dark"] section[data-testid="stSidebar"] ~ div .stChatInput {
            background: rgb(14, 17, 23);
            border-top: 1px solid #333;
        }
        
        /* Adjust width to account for sidebar */
        @media (min-width: 768px) {
            section[data-testid="stSidebar"] ~ div .stChatInput {
                left: var(--sidebar-width, 21rem);
                right: 0;
            }
        }
        
        /* When sidebar is collapsed */
        section[data-testid="stSidebar"][aria-expanded="false"] ~ div .stChatInput {
            left: 0;
        }
        </style>
        
        <script>
        // Add keyboard shortcuts support
        document.addEventListener('DOMContentLoaded', function() {
            // Find chat input field
            const observer = new MutationObserver(function(mutations) {
                const chatInput = document.querySelector('textarea[data-testid="stChatInput"]');
                if (chatInput && !chatInput.hasAttribute('data-shortcut-attached')) {
                    chatInput.setAttribute('data-shortcut-attached', 'true');
                    
                    // Add keyboard event listener
                    chatInput.addEventListener('keydown', function(e) {
                        // Enter (without Shift) - send message
                        if (e.key === 'Enter' && !e.shiftKey) {
                            e.preventDefault();
                            // Trigger the send button
                            const sendButton = document.querySelector('button[kind="primary"]');
                            if (sendButton) {
                                sendButton.click();
                            }
                        }
                        // Ctrl+Enter or Cmd+Enter - send message (alternative)
                        else if (e.key === 'Enter' && (e.ctrlKey || e.metaKey)) {
                            e.preventDefault();
                            const sendButton = document.querySelector('button[kind="primary"]');
                            if (sendButton) {
                                sendButton.click();
                            }
                        }
                        // Shift+Enter - new line (default behavior)
                    });
                }
            });
            
            observer.observe(document.body, {
                childList: true,
                subtree: true
            });
        });
        </script>
    """, unsafe_allow_html=True)
    
    # Check if we have a current chat
    if not st.session_state.current_dialogue_id:
        # Show welcome screen
        st.title("💬 Чат с RAG системой")
        st.markdown("---")
        
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            st.info("👋 Добро пожаловать! Создайте новый чат или выберите существующий из списка слева.")
            
            if st.button("🆕 Начать новый чат", type="primary", use_container_width=True):
                create_new_chat()
        
        return
    
    # Get cached messages
    current_messages = get_current_chat_messages()
    
    # Display chat header
    if current_messages:
        chat_name = get_chat_display_name(
            st.session_state.current_dialogue_id,
            current_messages[0]["query"]
        )
    else:
        chat_name = "Новый диалог"
    
    col1, col2 = st.columns([4, 1])
    with col1:
        st.title(f"💬 {chat_name}")
    with col2:
        if st.button("🗑️ Удалить чат", use_container_width=True):
            if delete_chat(st.session_state.current_dialogue_id):
                st.success("✅ Чат удален")
    
    st.markdown("---")
    
    # Chat messages container - load from DB
    if not current_messages:
        st.info("📝 Начните диалог, задав первый вопрос ниже")
    else:
        # Display all messages
        for msg in current_messages:
            # User message
            with st.chat_message("user"):
                st.markdown(msg["query"])
                timestamp_str = msg.get("timestamp", "")
                try:
                    dt = datetime.fromisoformat(timestamp_str)
                    st.caption(f"🕐 {dt.strftime('%H:%M:%S')}")
                except:
                    pass
            
            # Assistant message
            with st.chat_message("assistant"):
                st.markdown(msg["answer"])
                
                # Show reasoning in expander
                if msg.get("reason"):
                    with st.expander("📝 Обоснование"):
                        st.markdown(msg["reason"])
    
    # Input area - fixed at the bottom via CSS
    query = st.chat_input(
        "Введите ваш вопрос...",
        key="chat_input"
    )
    
    if query:
        # Send message and get response
        with st.spinner("🤔 Думаю..."):
            result = send_message(query)



# --- Main App ---
def main():
    st.set_page_config(
        page_title="RAG Chat System",
        page_icon="💬",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Initialize session state FIRST (before any other operations)
    init_session_state()
    
    # Initialize needs_rerun flag if not exists
    if "needs_rerun" not in st.session_state:
        st.session_state.needs_rerun = False
    
    # Initialize history table once using cache
    init_db()
    
    # Load chats list if not loaded yet
    if not st.session_state.chat_list_loaded:
        load_chats_list()
    
    # Sidebar
    with st.sidebar:
        st.title("💬 RAG Chat")
        
        # New chat button
        if st.button("➕ Новый чат", use_container_width=True, type="primary"):
            create_new_chat()
        
        st.markdown("---")
        
        # Chats list - use cached
        col1, col2 = st.columns([3, 1])
        with col1:
            st.subheader("📝 Ваши чаты")
        with col2:
            if st.button("🔄", help="Обновить список чатов"):
                st.session_state.chat_list_loaded = False
                load_chats_list()
        
        if not st.session_state.chat_list:
            st.info("Нет чатов. Создайте новый!")
        else:
            # Display chats from cache
            for chat in st.session_state.chat_list:
                dialogue_id = chat["dialogue_id"]
                message_count = chat.get("message_count", 0)
                started_at = chat.get("started_at", "")
                
                # Get chat name (only load history if chat has messages)
                if message_count > 0:
                    history = get_history_by_dialogue(dialogue_id)
                    first_query = history[0]["query"] if history else None
                else:
                    first_query = None
                chat_name = get_chat_display_name(dialogue_id, first_query)
                
                # Format time
                try:
                    dt = datetime.fromisoformat(started_at)
                    time_str = dt.strftime('%d.%m %H:%M')
                except:
                    time_str = ""
                
                # Check if this is current chat
                is_current = dialogue_id == st.session_state.current_dialogue_id
                
                # Format button text with chat name and metadata
                button_text = f"{'📌' if is_current else '💬'} {chat_name}\n💬 {message_count} • {time_str}"
                
                if st.button(
                    button_text,
                    key=f"chat_{dialogue_id}",
                    use_container_width=True,
                    type="primary" if is_current else "secondary"
                ):
                    switch_to_chat(dialogue_id)
    
    # Handle rerun at the end if needed
    if st.session_state.needs_rerun:
        st.session_state.needs_rerun = False
        st.rerun()
    
    # Main content area
    page_chat()


if __name__ == "__main__":
    main()