# ------------------- MUST BE FIRST -------------------
import streamlit as st
from pathlib import Path

# Create folder if it doesn't exist
KNOWLEDGE_DIR = Path("knowledge_base")
KNOWLEDGE_DIR.mkdir(parents=True, exist_ok=True)

st.set_page_config(
    page_title="Sirraya xBrain - Intelligent Assistant", 
    layout="centered",
    page_icon="🧠",
    initial_sidebar_state="collapsed"
)
# -----------------------------------------------------

from knowledge_engine import KnowledgeManager, Config


def initialize_lisa():
    """Initialize LISA knowledge manager"""
    if "lisa" not in st.session_state:
        with st.spinner("🔄 Initializing knowledge engine..."):
            try:
                st.session_state.lisa = KnowledgeManager()
                if st.session_state.lisa.qa_chain:
                    st.success("✅ Knowledge engine initialized successfully!")
                else:
                    st.error("❌ Failed to initialize knowledge engine. Please check your setup.")
            except Exception as e:
                st.error(f"❌ Error initializing system: {e}")
                st.session_state.lisa = None


def render_sidebar():
    """Render the sidebar for knowledge management"""
    with st.sidebar:
        st.header("📚 Knowledge Management")
        
        # File upload section
        uploaded_file = st.file_uploader(
            "Add knowledge file", 
            type=["txt"],
            help="Upload text files to expand LISA's knowledge base"
        )
        
        if uploaded_file:
            if st.session_state.lisa:
                save_path = KNOWLEDGE_DIR / uploaded_file.name
                try:
                    # Save the uploaded file into knowledge_base folder
                    with open(save_path, "wb") as f:
                        f.write(uploaded_file.getbuffer())
                    st.success(f"✅ Saved {uploaded_file.name} to knowledge_base folder")
                    st.info("💡 Click 'Rebuild Knowledge Base' to update the index")
                except Exception as e:
                    st.error(f"❌ Error saving {uploaded_file.name}: {e}")
            else:
                st.error("❌ Knowledge engine not initialized")
        
        # Rebuild button
        if st.button("🔄 Rebuild Knowledge Base", type="primary"):
            with st.spinner("🔧 Rebuilding knowledge engine..."):
                try:
                    st.session_state.lisa = KnowledgeManager()
                    if st.session_state.lisa.qa_chain:
                        st.success("✅ Knowledge base rebuilt successfully!")
                        st.experimental_rerun()
                    else:
                        st.error("❌ Failed to rebuild knowledge base")
                except Exception as e:
                    st.error(f"❌ Error rebuilding: {e}")
        
        st.divider()
        
        # System info section
        st.subheader("🔧 System Info")
        st.info("**Embedding Model:** `mxbai-embed-large`")
        st.info("**LLM Model:** `phi`")
        st.info("**Retrieval:** Hybrid (Vector + BM25)")
        
        # Knowledge base stats
        if st.session_state.lisa:
            file_count = st.session_state.lisa.get_knowledge_files_count()
            st.metric("📄 Knowledge Files", file_count)


def render_chat_interface():
    """Render the main chat interface"""
    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    # Display chat history
    for msg in st.session_state.messages:
        with st.chat_message(msg["role"]):
            st.write(msg["content"])
            if msg["role"] == "assistant" and msg.get("sources"):
                with st.expander("📖 View Sources", expanded=False):
                    for i, source in enumerate(msg["sources"]):
                        st.markdown(f"**📄 Source {i+1}:**")
                        st.text(source.page_content[:300] + "..." if len(source.page_content) > 300 else source.page_content)
                        if hasattr(source, 'metadata') and source.metadata:
                            st.caption(f"From: {source.metadata.get('source', 'Unknown')}")
    
    # Handle new user query
    if prompt := st.chat_input("Ask LISA about anything in the knowledge base..."):
        # Add user message
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        with st.chat_message("user"):
            st.write(prompt)
        
        # Generate response
        with st.chat_message("assistant"):
            if st.session_state.lisa and st.session_state.lisa.qa_chain:
                with st.spinner("🤔 Thinking..."):
                    result = st.session_state.lisa.query(prompt)
                    
                    st.write(result["answer"])
                    
                    # Show processing time
                    if result["processing_time"] > 0:
                        st.caption(f"⚡ Processed in {result['processing_time']:.0f}ms")
                    
                    # Store message with sources
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": result["answer"],
                        "sources": result["source_chunks"] if result["source_chunks"] else None
                    })
            else:
                error_msg = "❌ LISA is not properly initialized. Please try rebuilding the knowledge base."
                st.error(error_msg)
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": error_msg
                })


def main():
    """Main application function"""
    # Header
    st.title("🧠 Sirraya xBrain - LISA")
    st.markdown("*Intelligent Assistant powered by Advanced RAG Technology*")
    
    # Initialize LISA
    initialize_lisa()
    
    # Render sidebar
    render_sidebar()
    
    # Render chat interface
    render_chat_interface()


if __name__ == "__main__":
    main()