import openai
import gradio as gr
import os

openai.api_key = os.getenv("OPENAI_API_KEY")
MODEL_NAME = "gpt-4"

# Function to send messages to GPT-4
global_chat_history = ""

def generate_response(prompt):
  global global_chat_history
  # 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)],
    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()