pages/document_store.py

# pages/document_store.py
import streamlit as st
from utils.database import (
    create_collection,
    get_collections,
    get_collection_documents,
    delete_collection,
    add_document_to_collection,
    process_and_store_document
)

def display_collection_manager():
    """Display the collection management interface."""
    st.header("Collection Management")
    
    # Get database connection from session state
    conn = st.session_state.db_conn
    
    # Create new collection section
    with st.expander("Create New Collection", expanded=False):
        col1, col2 = st.columns([3, 1])
        with col1:
            name = st.text_input("Collection Name", key="new_collection_name")
            description = st.text_area("Description", key="new_collection_desc")
        with col2:
            if st.button("Create Collection", key="create_collection_btn") and name:
                if create_collection(conn, name, description):
                    st.success(f"Collection '{name}' created!")
                    st.rerun()

    # Display existing collections
    collections = get_collections(conn)
    if collections:
        st.subheader("Existing Collections")
        
        for collection in collections:
            with st.expander(f"📁 {collection['name']} ({collection['doc_count']} documents)"):
                # Collection details
                st.write(f"**Description:** {collection['description']}")
                st.write(f"**Created:** {collection['created_at']}")
                
                # Get documents in this collection
                documents = get_collection_documents(conn, collection['id'])
                
                # Display documents
                if documents:
                    st.write("**Documents:**")
                    for doc in documents:
                        st.write(f"• {doc['name']}")
                
                # Collection management buttons
                col1, col2, col3 = st.columns([1, 1, 1])
                with col1:
                    if st.button("Add Documents", key=f"add_{collection['id']}"):
                        st.session_state.selected_collection = collection['id']
                        st.session_state.show_upload = True
                with col2:
                    if st.button("Edit", key=f"edit_{collection['id']}"):
                        st.session_state.editing_collection = collection['id']
                with col3:
                    if st.button("Delete", key=f"del_{collection['id']}"):
                        if delete_collection(conn, collection['id']):
                            st.success(f"Collection '{collection['name']}' deleted!")
                            st.rerun()

def handle_document_upload_with_collection(uploaded_files, collection_id):
    """Handle document upload with collection assignment."""
    conn = st.session_state.db_conn
    
    for uploaded_file in uploaded_files:
        # Process and store the document
        document_id = process_and_store_document(uploaded_file)
        
        if document_id:
            # Add to collection
            if add_document_to_collection(conn, document_id, collection_id):
                st.success(f"Added {uploaded_file.name} to collection!")
            else:
                st.error(f"Failed to add {uploaded_file.name} to collection!")

# Initialize required session state
if 'show_upload' not in st.session_state:
    st.session_state.show_upload = False
if 'selected_collection' not in st.session_state:
    st.session_state.selected_collection = None
if 'editing_collection' not in st.session_state:
    st.session_state.editing_collection = None

# Add this if you need it for the page
def main():
    display_collection_manager()

if __name__ == "__main__":
    main()