import gradio as gr
import openai
import os
import json

openai.api_key = os.environ["OPENAI_API_KEY"]
openai.base_url = "https://api.ozone-ai.com/v1/"

log_file = "gpt4_chat_logs.jsonl"

def log_message(messages):
    with open(log_file, "a") as f:
        json.dump(messages, f)
        f.write("\n")

def respond(message, chat_history):
    messages = [{"role": "system", "content": "You are a helpful assistant."}]
    for user_msg, bot_msg in chat_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    bot_message = ""
    try:
        chat_history.append((message, ""))
        for response in openai.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            max_tokens=5000,
            temperature=0.7,
            top_p=1,
            frequency_penalty=0,
            presence_penalty=0,
            stream=True
        ):
            bot_message += (response.choices[0].delta.content or "")
            chat_history[-1] = (message, bot_message)
            yield "", chat_history
    except Exception as e:
        print(f"Error during API call: {e}")
        bot_message = "Sorry, I encountered an error. Please try again later."
        if chat_history:
            chat_history[-1] = (message, bot_message)
        else:
            chat_history.append((message, bot_message))
            
    log_messages = [{"role": "system", "content": "You are a helpful assistant."}]
    for user_msg, bot_msg in chat_history:
        log_messages.append({"role": "user", "content": user_msg})
        log_messages.append({"role": "assistant", "content": bot_msg})

    log_message(log_messages)
    
    return "", chat_history


with gr.Blocks() as demo:
    gr.Markdown("")

    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    def clear_chat():
        return None

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(clear_chat, None, chatbot, queue=False)

demo.launch()