import streamlit as st
from doc_manager import DocumentManager, ChatInterface

def main():
    st.set_page_config(page_title="RFP Analyzer", layout="wide")
    
    # Initialize document manager
    if 'doc_manager' not in st.session_state:
        st.session_state.doc_manager = DocumentManager()
    
    # Sidebar for document management
    with st.sidebar:
        st.title("📁 Document Management")
        
        # Collection management
        st.subheader("Collections")
        
        # Create new collection
        new_collection = st.text_input("New Collection Name")
        if st.button("Create Collection") and new_collection:
            collection_id = st.session_state.doc_manager.create_collection(new_collection)
            st.success(f"Created collection: {new_collection}")
            st.rerun()
        
        # Collection selection
        collections = st.session_state.doc_manager.get_collections()
        collection_options = ["All Documents"] + [c['name'] for c in collections]
        selected_collection = st.selectbox("Select Collection", collection_options)
        
        # Document upload
        st.subheader("Upload Documents")
        uploaded_files = st.file_uploader(
            "Upload PDF documents",
            type=['pdf'],
            accept_multiple_files=True
        )
        
        if uploaded_files:
            collection_id = None
            if selected_collection != "All Documents":
                collection_id = next(c['id'] for c in collections if c['name'] == selected_collection)
            
            doc_ids = st.session_state.doc_manager.upload_documents(uploaded_files, collection_id)
            if doc_ids:
                st.success(f"Uploaded {len(doc_ids)} documents")
                st.rerun()
    
    # Main content area
    st.title("RFP Analyzer")
    
    # Document selection for chat
    if selected_collection != "All Documents":
        collection_id = next(c['id'] for c in collections if c['name'] == selected_collection)
        documents = st.session_state.doc_manager.get_collection_documents(collection_id)
    else:
        documents = st.session_state.doc_manager.get_collection_documents()
    
    if documents:
        st.subheader("Available Documents")
        selected_docs = st.multiselect(
            "Select documents for analysis",
            options=[doc['id'] for doc in documents],
            format_func=lambda x: next(doc['name'] for doc in documents if doc['id'] == x)
        )
        
        if selected_docs:
            if st.button("Start Chat"):
                vector_store = st.session_state.doc_manager.initialize_chat(selected_docs)
                if vector_store:
                    st.session_state.chat_interface = ChatInterface(vector_store)
                    st.session_state.chat_active = True
                    st.rerun()
    
    # Chat interface
    if st.session_state.get('chat_active'):
        st.divider()
        st.subheader("Chat Interface")
        st.session_state.chat_interface.display()

if __name__ == "__main__":
    main()