import openai
import gradio as gr
import os

MODEL_NAME = "gpt-4"
global_chat_history = ""


def generate_response(prompt, api_key):
    global global_chat_history
    openai.api_key = api_key
    # Send message to GPT-4
    message = {"role": "user", "content": prompt}
    chat_history_messages = [{"role": "user", "content": m}
                             for m in global_chat_history.split("\n")] + [message]
    response = openai.ChatCompletion.create(
        model=MODEL_NAME,
        messages=chat_history_messages,
        temperature=1,
        max_tokens=1024,
        n=1,
        stop=None,
    )
    # Update global chat history with the new message and response
    global_chat_history += f"\nUser: {prompt}\nAI: {response.choices[0].message.content}"
    return response.choices[0].message.content


blocks = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(placeholder="Enter your prompt...", lines=2),
        gr.Textbox(label="API Key", type="password")
    ],
    outputs=[gr.Textbox(placeholder="Response", lines=10, readonly=True)],
    title="GPT-4 Chatbot",
    theme="compact",
    layout="vertical",
    allow_flagging=False,
)


def add_cors_headers(response):
    response.headers["Access-Control-Allow-Origin"] = "*"
    response.headers["Access-Control-Allow-Headers"] = "Content-Type,Authorization"
    response.headers["Access-Control-Allow-Methods"] = "POST,OPTIONS"
    return response

# @blocks.flask_app.after_request


def enable_cors(response):
    return add_cors_headers(response)


def save_conversation():
    global global_chat_history
    if global_chat_history is None:
        return ""
    with open('chat_history.txt', 'a') as f:
        f.write(global_chat_history)
        f.write("\n")
    return ""


with blocks:
    b1 = gr.Button("Save Conversation")
    b1.click(save_conversation)

blocks.launch()