# src/components/knowledge_base.py
import streamlit as st
from datetime import datetime

def handle_doc_select():
    st.session_state.current_chat = True
    st.session_state.chat_history = []

def handle_start_chat():
    st.session_state.current_chat = True
    st.session_state.chat_history = []

def display_knowledge_base(conn, backend):
    st.markdown("### 📚 Knowledge Base")
    
    if conn is not None:
        try:
            cursor = conn.cursor()
            cursor.execute("SELECT id, name, upload_date FROM documents ORDER BY upload_date DESC")
            documents_in_db = cursor.fetchall()
            
            if documents_in_db:
                st.markdown("#### Available Documents")
                
                for doc_id, name, upload_date in documents_in_db:
                    col1, col2 = st.columns([3, 1])
                    
                    with col1:
                        selected = st.checkbox(
                            name,
                            value=doc_id in st.session_state.selected_docs,
                            key=f"doc_{doc_id}",
                            on_change=handle_doc_select
                        )
                        
                        if selected and doc_id not in st.session_state.selected_docs:
                            st.session_state.selected_docs.append(doc_id)
                        elif not selected and doc_id in st.session_state.selected_docs:
                            st.session_state.selected_docs.remove(doc_id)
                    
                    with col2:
                        upload_date = datetime.strptime(upload_date, '%Y-%m-%d %H:%M:%S')
                        st.text(upload_date.strftime('%Y-%m-%d'))
                
                initialize_selected_documents(conn, backend)
                
                if st.session_state.selected_docs:
                    st.button("🚀 Start New Chat", 
                             on_click=handle_start_chat,
                             use_container_width=True)
            else:
                st.info("No documents in the knowledge base. Upload some documents to get started!")
                
        except Exception as e:
            st.error(f"Error accessing knowledge base: {e}")

def initialize_selected_documents(conn, backend):
    if st.session_state.selected_docs and not st.session_state.documents_initialized:
        with st.spinner("Initializing document analysis..."):
            selected_documents = []
            selected_doc_names = []
            
            cursor = conn.cursor()
            for doc_id in st.session_state.selected_docs:
                cursor.execute(
                    "SELECT content, name FROM documents WHERE id = ?",
                    (doc_id,)
                )
                result = cursor.fetchone()
                if result:
                    selected_documents.append(result[0])
                    selected_doc_names.append(result[1])
            
            embeddings = backend.get_embeddings_model()
            if embeddings:
                vector_store = backend.initialize_faiss(
                    embeddings,
                    selected_documents,
                    selected_doc_names
                )
                if vector_store:
                    st.session_state.qa_system = backend.initialize_qa_system(vector_store)
                    st.session_state.documents_initialized = True