# components/collection_manager.py

import streamlit as st  # Import Streamlit here

from utils.database import (
    get_all_documents,
    create_collection,
    add_document_to_collection,
    get_collections,
    search_documents,
    get_collection_documents,
)
st.session_state.reinitialize_chat = True
def show_collection_creation_dialog():
    """Show dialog for creating a new collection."""
    with st.form("create_collection"):
        st.subheader("Create New Collection")
        name = st.text_input("Collection Name", key="new_col_name")
        description = st.text_area("Description", key="new_col_desc")
        
        # Document selection
        existing_docs = get_all_documents(st.session_state.db_conn)
        if existing_docs:
            st.subheader("Add Existing Documents")
            selected_docs = st.multiselect(
                "Select documents to add",
                options=[doc['id'] for doc in existing_docs],
                format_func=lambda x: next(doc['name'] for doc in existing_docs if doc['id'] == x)
            )
        
        submitted = st.form_submit_button("Create Collection")
        
        if submitted and name:
            collection_id = create_collection(st.session_state.db_conn, name, description)
            if collection_id and selected_docs:
                for doc_id in selected_docs:
                    add_document_to_collection(st.session_state.db_conn, doc_id, collection_id)
            st.success(f"Collection '{name}' created!")
            st.session_state.show_collection_dialog = False
            st.rerun()

def enhanced_document_preview(doc):
    """Enhanced document preview with actions."""
    col1, col2 = st.columns([3, 1])
    
    with col1:
        st.text_area(
            "Preview",
            value=doc['content'][:1000] + "..." if len(doc['content']) > 1000 else doc['content'],
            height=200,
            disabled=True
        )
    
    with col2:
        st.button("🗑️ Remove", key=f"remove_{doc['id']}", use_container_width=True)
        st.button("📋 Copy", key=f"copy_{doc['id']}", use_container_width=True)
        st.button("💬 Chat", key=f"chat_{doc['id']}", use_container_width=True)

def display_enhanced_collections():
    """Enhanced collections management interface."""
    col1, col2 = st.columns([1, 2])
    
    with col1:
        st.subheader("Collections")
        
        # Search collections
        search_query = st.text_input("🔍 Search collections", key="collection_search")
        
        # Create new collection button
        if st.button("+ New Collection", type="primary", use_container_width=True):
            st.session_state.show_collection_dialog = True
        
        # Collection list with stats
        collections = get_collections(st.session_state.db_conn)
        if search_query:
            collections = [c for c in collections 
                         if search_query.lower() in c['name'].lower()]
        
        st.divider()
        
        for collection in collections:
            with st.container():
                col_button = st.button(
                    f"📁 {collection['name']}",
                    key=f"col_{collection['id']}",
                    use_container_width=True
                )
                
                # Show collection stats
                stats_cols = st.columns(3)
                with stats_cols[0]:
                    st.caption(f"📄 {collection['doc_count']}")
                with stats_cols[1]:
                    st.caption(f"📅 {collection['created_at'][:10]}")
                with stats_cols[2]:
                    st.caption("⚙️ Manage")
                
                if col_button:
                    st.session_state.selected_collection = collection
                    st.rerun()
    
    with col2:
        if st.session_state.get('selected_collection'):
            display_collection_content(st.session_state.selected_collection)
        else:
            st.info("Select a collection to view its contents")

def display_collection_content(collection):
    """Display enhanced collection contents."""
    st.subheader(collection['name'])
    
    # Collection actions
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.button("📥 Add Documents", use_container_width=True)
    with col2:
        st.button("✏️ Edit Collection", use_container_width=True)
    with col3:
        st.button("📤 Export", use_container_width=True)
    with col4:
        st.button("🗑️ Delete", use_container_width=True)
    
    # Document search within collection
    search_query = st.text_input("🔍 Search documents in collection")
    
    # Filters
    with st.expander("Filters"):
        filter_cols = st.columns(3)
        with filter_cols[0]:
            st.date_input("Date Range")
        with filter_cols[1]:
            st.selectbox("Sort By", ["Name", "Date", "Size"])
        with filter_cols[2]:
            st.selectbox("View", ["List", "Grid"])
    
    # Document list with enhanced display
    documents = get_collection_documents(st.session_state.db_conn, collection['id'])
    if search_query:
        documents = [doc for doc in documents 
                    if search_query.lower() in doc['name'].lower() or 
                       search_query.lower() in doc.get('content', '').lower()]
    
    # Display documents in a clean, modern format
    for doc in documents:
        with st.container():
            doc_cols = st.columns([0.1, 2, 0.5])
            with doc_cols[0]:
                st.checkbox("", key=f"select_{doc['id']}")
            with doc_cols[1]:
                st.markdown(f"**{doc['name']}**")
                st.caption(f"Uploaded: {doc.get('upload_date', 'N/A')}")
            with doc_cols[2]:
                st.button("Preview", key=f"preview_{doc['id']}", use_container_width=True)
                st.button("Remove", key=f"remove_{doc['id']}", use_container_width=True)

        # Document preview expander
        if st.session_state.get(f"preview_{doc['id']}", False):
            with st.expander("Document Preview", expanded=True):
                st.text_area(
                    "Content",
                    value=doc['content'][:1000] + "..." if len(doc['content']) > 1000 else doc['content'],
                    height=200,
                    disabled=True
                )

def show_document_preview(document):
    """Show document preview in a modal."""
    with st.expander(f"Preview: {document['name']}", expanded=True):
        st.text_area(
            "Content",
            value=document['content'][:1000] + "..." if len(document['content']) > 1000 else document['content'],
            height=300,
            disabled=True
        )