import streamlit as st
from helper import load_hf_datasets, search, get_file_paths, get_images_from_s3_to_display
import os


#Load environment variables
AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
# Predefined list of datasets
datasets = ["WayveScenes", "StopSign_test"]  # Example dataset names

# AWS S3 bucket name
bucket_name = "datasets-quasara-io"

# Streamlit App
def main():
    st.title("Semantic Search and Image Display")

    # Select dataset from dropdown
    dataset_name = st.selectbox("Select Dataset", datasets)
    if dataset_name == 'WayveScenes':
        folder_path = 'WayveScenes/'
    else:
        folder_path = ''
    # Input search query
    query = st.text_input("Enter your search query")

    # Number of results to display
    limit = st.number_input("Number of results to display", min_value=1, max_value=10, value=10)

    # Search button
    if st.button("Search"):
        # Validate input
        if not query:
            st.warning("Please enter a search query.")
        else:
            # Load the selected dataset
            df = load_hf_datasets(dataset_name)

            # Perform the search
            results = search(query, df, limit, 0, "cosine", search_in_images=True, search_in_small_objects=False)

            # Get the S3 file paths of the top results
            top_k_paths = get_file_paths(df, results)

            # Display images from S3
            if top_k_paths:
                st.write(f"Displaying top {len(top_k_paths)} results for query '{query}':")

                get_images_from_s3_to_display(bucket_name, top_k_paths,AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, folder_path)
            else:
                st.write("No results found.")

if __name__ == "__main__":
    main()


#the query and the stuff is not really correlating