import streamlit as st
from SearchEngine import searchengine

def main():
    st.title("ChromaDB Search Engine")

    model_name = "your_model_name"  # Replace with the actual model name
    collection_name = "your_collection_name"  # Replace with the actual collection name

    search_engine = searchengine(model_name, collection_name)

    st.sidebar.header("Add Document")
    text_input = st.sidebar.text_area("Enter Document Text")
    metadata_input = st.sidebar.text_input("Enter Metadata")
    add_button = st.sidebar.button("Add Document")

    if add_button:
        document_id = search_engine.count() + 1  # Incremental ID
        search_engine.add(text_input, metadata_input, document_id)
        st.sidebar.success(f"Document added with ID: {document_id}")

    st.sidebar.header("Search")
    query = st.sidebar.text_input("Enter Search Query")
    search_button = st.sidebar.button("Search")

    if search_button:
        results = search_engine.query(query)
        st.subheader("Search Results:")
        for result in results:
            st.write(f"Document ID: {result['id']}, Metadata: {result['metadata']}")
            st.write(f"Text: {result['text']}")
            st.markdown("---")

if __name__ == "__main__":
    main()