app.py

import streamlit as st
import os
import time
from src.retrieval import RetrievalEngine

# --- PAGE CONFIGURATION ---
st.set_page_config(
    page_title="Pro RAG Enterprise",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# --- CUSTOM CSS ---
st.markdown("""
<style>
    .stChatInputContainer {
        padding-bottom: 20px;
    }
    .block-container {
        padding-top: 30px;
    }
    h1 {
        color: #0F172A;
    }
    .stSidebar {
        background-color: #F8FAFC;
        border-right: 1px solid #E2E8F0;
    }
    /* Status Badge Style */
    .status-badge {
        padding: 4px 8px;
        border-radius: 4px;
        font-size: 0.8em;
        font-weight: bold;
    }
</style>
""", unsafe_allow_html=True)

# --- 1. INITIALIZE ENGINE (Cached) ---
@st.cache_resource
def get_engine():
    return RetrievalEngine()

# Initialize and Check Connection Type
try:
    engine = get_engine()
    
    # Check env vars to see where we are connected
    if os.getenv("QDRANT_URL"):
        conn_type = "☁️ Qdrant Cloud"
        status_color = "green"
    else:
        conn_type = "🏠 Local Docker"
        status_color = "orange"
        
    db_status = f"{conn_type} Connected"
    
except Exception as e:
    engine = None
    db_status = f"❌ Error: {e}"
    status_color = "red"

# --- 2. SIDEBAR (The Control Panel) ---
with st.sidebar:
    st.title("🎛️ Control Panel")
    
    # Connection Status
    st.markdown(f"**System Status:** :{status_color}[{db_status}]")
    st.divider()

    # Mode Selection
    st.subheader("🔍 Search Mode")
    mode_display = {
        "Global Search (All Data)": "all",
        "📄 PDF Documents (Financials)": "pdf",
        "📊 Structured Data (Excel/CSV)": "csv",
        "🖼️ Visual Intelligence (Graphs)": "visual"
    }
    
    selected_mode_label = st.selectbox(
        "Select Knowledge Source:",
        list(mode_display.keys()),
        index=0
    )
    # Convert label back to backend keyword
    filter_mode = mode_display[selected_mode_label]

    st.info(
        f"""
        **Current Focus:** {selected_mode_label}
        
        *Engine filters retrieval to strictly match this data type.*
        """
    )
    
    st.divider()
    if st.button("🗑️ Clear Chat History"):
        st.session_state.messages = []
        st.rerun()

# --- 3. MAIN CHAT INTERFACE ---

st.title("🤖 Enterprise Knowledge Assistant")
st.caption("Level 1 Pro RAG System | Powered by Qdrant & GPT-4o")

# Initialize Chat History
if "messages" not in st.session_state:
    st.session_state.messages = []

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

# --- 4. HANDLE USER INPUT ---
if prompt := st.chat_input("Ask a question about your data..."):
    # A. Display User Message
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # B. Generate AI Response
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        
        with st.spinner(f"Searching {selected_mode_label}..."):
            try:
                # CALL THE BACKEND
                response_text = engine.query(prompt, filter_type=filter_mode)
                
                # Display response
                message_placeholder.markdown(response_text)
                
            except Exception as e:
                error_msg = f"❌ System Error: {str(e)}"
                message_placeholder.error(error_msg)
                response_text = error_msg

    # C. Save AI Message
    st.session_state.messages.append({"role": "assistant", "content": response_text})