from openai import OpenAI
import os
import time
import gradio as gr

# Initialize clients with API keys
client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", "<your OpenAI API key if not set as env var>"))

# Step 1: Create an Assistant
assistant = client.beta.assistants.create(
    name="WEA Emergency Message Generator",
    instructions="You are an emergency message generator. You will write emergency messages of 360 characters or less to broadcast to wireless phones. Do not include hashtags or emojis.",
    model="gpt-4-turbo-preview",
    tools=[{"type": "retrieval"}]
)

# Step 2: Create a Thread
thread = client.beta.threads.create()

def main(query):
    # Step 3: Add a Message to a Thread
    message = client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=query
    )

    # Step 4: Run the Assistant
    run = client.beta.threads.runs.create(
        thread_id=thread.id,
        assistant_id=assistant.id,
        instructions="You are an emergency message generator. You will write emergency messages of 360 characters or less to broadcast to wireless phones. Do not include hashtags or emojis."
    )

    while True:
        # Wait for 5 seconds
        time.sleep(5)

        # Retrieve the run status
        run_status = client.beta.threads.runs.retrieve(
            thread_id=thread.id,
            run_id=run.id
        )

        # If run is completed, get messages
        if run_status.status == 'completed':
            messages = client.beta.threads.messages.list(
                thread_id=thread.id,
		limit=1
            )
            response = ""
            # Loop through messages and print content based on role
            for msg in messages.data:
                role = msg.role
                content = msg.content[0].text.value
                response += f"{role.capitalize()}: {content}\n\n"
            return response+"\n\n"
        else:
            continue

# Create a Gradio Interface
iface = gr.Interface(
	fn=main, 
	inputs=[gr.Textbox(label="Describe the warning here", lines=4)], 
	outputs=[gr.Textbox(label="Suggested Message", lines=4)], 
	allow_flagging="never",
	title="WEA Emergency Message Assistant", 
	description="Hello. I'm Emmy your emergency messaging bot. Please describe the message you would like me to generate.",
	article="All input and output will be saved for research purposes.").launch()
iface.launch()