import streamlit as st
from SearchEngine import searchengine
import logging


# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
model_name = "all-MiniLM-L6-v2" 
collection_name = "docs"  

search_engine = searchengine(model_name, collection_name)


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

    st.sidebar.header("Add Document")
    text_input = st.sidebar.text_area("Text")
    metadata_input = {'type':st.sidebar.text_input("Type")}
    add_button = st.sidebar.button("Save")

    if add_button:
        document_id = str(search_engine.count() + 1)
        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("Query")
    search_button = st.sidebar.button("Search")

    if search_button:
        results = search_engine.query(query)
        st.subheader("Search Results:")
        logging.info("result :")
        logging.info(results)
        ids = results['ids'][0]
        distances = results['distances'][0]
        metadatas = results['metadatas'][0]
        documents = results['documents'][0]

        for index, id in enumerate(ids):
            
            st.write(f"Document ID: {id}, Metadata: {metadatas[index]}")
            st.write(f"Text: {documents[index]}")
            st.write(f"distance: {distances[index]}")
            st.markdown("---")

if __name__ == "__main__":
    main()