import streamlit as st

def display_initial_buttons():
    if "upload_flow" not in st.session_state:
        st.session_state.upload_flow = False
    if "query_flow" not in st.session_state:
        st.session_state.query_flow = False

    if st.button("Upload new CSV"):
        st.session_state.upload_flow = True
        st.session_state.query_flow = False

    if st.button("Query existing data"):
        st.session_state.query_flow = True
        st.session_state.upload_flow = False

def display_class_dropdown(client):
    if st.session_state.upload_flow:
        existing_classes = [cls["class"] for cls in client.schema.get()["classes"]]
        class_options = existing_classes + ["New Class"]
        return st.selectbox("Select a class or create a new one:", class_options)
    elif st.session_state.query_flow:
        existing_classes = [cls["class"] for cls in client.schema.get()["classes"]]
        class_options = existing_classes + ["Query all data"]
        return st.selectbox("Select a class or query all data:", class_options)

def handle_new_class_selection(client, selected_class):
    if selected_class == "New Class":
        class_name = st.text_input("Enter the new class name:")
        class_description = st.text_input("Enter a description for the class:")
        if class_name and class_description:
            if st.button("Create Vector DB Class"):
                # Call function to create new class schema in Weaviate
                create_new_class_schema(class_name, class_description)

def csv_upload_and_ingestion(client, selected_class):
    csv_file = st.file_uploader("Upload a CSV file", type=["csv"])
    if csv_file:
        if st.button("Confirm CSV upload"):
            # Call function to ingest CSV data into Weaviate
            ingest_csv_to_weaviate(csv_file, selected_class)

def display_query_input():
    question = st.text_input("Enter your question:")
    if question:
        if st.button("Submit Query"):
            # Call function to query TAPAS with selected data and entered question
            query_tapas_with_weaviate_data(st.session_state.data_source, question)