import os
import openai
import gradio as gr
import time
from jtar import jtar_txt

# Set up password
username = os.getenv('username')
password = os.getenv('password')

# Set up OpenAI client
openai_api_key = os.getenv('openai_api_key')
openai.api_key = openai_api_key
client = openai.Client(api_key=openai.api_key)

# Set up assistant
assistant_id = os.getenv('assistant_id')
assistant = client.beta.assistants.retrieve(assistant_id)
thread = client.beta.threads.create()

# Get avatar link
avatar_link = os.getenv('avatar_link')

def chat_with_assistant(message, history):
    # Add the user's message to the thread
    client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=message
    )
    
    # Run the assistant
    run = client.beta.threads.runs.create(
        thread_id=thread.id,
        assistant_id=assistant_id
    )
    
    # Wait for the assistant's response
    while True:
        run_status = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
        if run_status.status == 'completed':
            # Retrieve the assistant's response
            messages = client.beta.threads.messages.list(thread_id=thread.id)
            assistant_response = messages.data[0].content[0].text.value
            break
        time.sleep(1)
    
    return assistant_response

# Custom CSS for chat bubbles and colors
custom_css = """
.user-message { background-color: #DCF8C6; }
.assistant-message { background-color: #E2E2E2; }
"""

# Create the Gradio interface
with gr.Blocks(css=custom_css) as demo:
    gr.HTML(f"""<img src = "https://i.postimg.cc/9MXZpwMy/a-captivating-pixar-style-3d-animation-featuring-a-ydzhw1k-R-GZi-Np1y7-Ufi-A-tvqb-XZ-z-R4a-G-8pf-Xcm-Og.png" width="320" height="240">""")
    gr.Markdown("# **NewsWave Companion** 🌊🐬")
    chatbot = gr.Chatbot(
        [],
        elem_id="chatbot",
        bubble_full_width=False,
        avatar_images=(None, avatar_link)
    )
    msg = gr.Textbox(
        show_label=False,
        placeholder="Enter text and press enter",
    )

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history):
        user_message = history[-1][0]
        bot_message = chat_with_assistant(user_message, history)
        history[-1][1] = bot_message
        return history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    gr.Markdown(jtar_txt)

demo.launch(auth=(username, password))