import openai
import time
import logging
import streamlit as st

openai_api_key = st.secrets["openai_api_key"]
client = openai.OpenAI(api_key = openai_api_key)
model = "gpt-3.5-turbo-16k"

# === Thread an empty thread
thread = client.beta.threads.create()
thread_id = thread.id

def wait_for_run_completion(client, thread_id, run_id, sleep_interval=5):
    """
    Waits for a run to complete and prints the elapsed time.:param client: The OpenAI client object.
    :param thread_id: The ID of the thread.
    :param run_id: The ID of the run.
    :param sleep_interval: Time in seconds to wait between checks.
    """
    while True:
        try:
            run = client.beta.threads.runs.retrieve(thread_id = thread_id, run_id = run_id)
            if run.completed_at:
                elapsed_time = run.completed_at - run.created_at
                formatted_elapsed_time = time.strftime(
                    "%H:%M:%S", time.gmtime(elapsed_time)
                )
                print(f"Run completed in {formatted_elapsed_time}")
                logging.info(f"Run completed in {formatted_elapsed_time}")
                # Get messages here once Run is completed!
                messages = client.beta.threads.messages.list(thread_id=thread_id)
                last_message = messages.data[0]
                response = last_message.content[0].text.value
                st.write(response)
                break
        except Exception as e:
            logging.error(f"An error occurred while retrieving the run: {e}")
            break
        logging.info("Waiting for run to complete...")
        time.sleep(sleep_interval)

def create_new_assistant():
    personal_trainer_assis = client.beta.assistants.create(
        name="Data Analyst",
        instructions="""You are great at creating beautiful data visualizations. You analyze data present in .csv files, understand trends, and come up with data visualizations relevant to those trends. You also share a brief text summary of the trends observed.""",
        model=model
    )
    print(f"Created new assistant with ID: {personal_trainer_assis.id}")
    return personal_trainer_assis.id

def main():
    # Streamlit interface
    st.title("Data Analyst")

    # Note that API key's running out of budget
    contact_url = "https://www.linkedin.com/in/linhvuu"
    st.write("If no result returns, it means I am running out of energy. Please contact [Linh Vuu](%s) to wake me up." % contact_url)
    
    uploaded_file = st.file_uploader("Upload a CSV file for analysis", type=['csv'])

    if uploaded_file is not None:

        file_content = uploaded_file.read()

        # Upload a file with an "assistants" purpose
        file_array = [file_content]
        file_id_array = []
        for file in file_array:
            file = client.files.create(
                file = file_content,
                purpose="assistants"
            )
            file_id_array.append(file.id)

        print(file_id_array)

    with st.form(key="user_input_form"):

        question = st.text_input("Enter question:")
        submit_button = st.form_submit_button(label="Response to my question above")

        if submit_button:

            # ==== Create a Message ====
            message = question
            message = client.beta.threads.messages.create(
                thread_id = thread_id, role = "user", content = message, file_ids = file_id_array
            )

            # Define assistant_id before the try-except block
            assistant_id = "asst_IY1VCWqTnQ2Uvj9zWHLXmblg"
            run = None

            # === Run our Assistant ===
            try:
                if assistant_id is None:
                    # If assistant_id is None, create a new assistant
                    assistant_id = create_new_assistant()

                run = client.beta.threads.runs.create(
                    thread_id = thread_id,
                    assistant_id = assistant_id
                )
            except openai.NotFoundError:
                # If assistant_id does not exist, create a new assistant and retry
                assistant_id = create_new_assistant()
                run = client.beta.threads.runs.create(
                    thread_id = thread_id,
                    assistant_id = assistant_id
                )

            # === Run ===
            wait_for_run_completion(client=client, thread_id=thread_id, run_id=run.id)

if __name__ == "__main__":
    main()