#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @时间: 2023/3/9 15:17
# @文件: chat_gpt_ui.py
import tkinter
from tkinter import simpledialog
import tkinter as tk
from tkinter import messagebox
 
import openai
import json
 
 
class ChatGPTUI:
 
    def __init__(self, root):
        self.root = root
        self.root.title('ChatGPT')
        self.root.geometry('500x600')
 
        # Set up the menu
        menubar = tk.Menu(self.root)
        self.root.config(menu=menubar)
        filemenu = tk.Menu(menubar, tearoff=0)
        filemenu.add_command(label='Exit', command=self.root.quit)
        menubar.add_cascade(label='File', menu=filemenu)
        editmenu = tk.Menu(menubar, tearoff=0)
        editmenu.add_command(label='Set API key', command=self.set_api_key)
        menubar.add_cascade(label='Edit', menu=editmenu)
 
        # Set up the chat window
        self.chat_window = tk.Text(self.root, state=tk.DISABLED, height=20, width=60)
        self.chat_window.pack(pady=10)
 
        # Set up the input field and button
        self.input_field = tk.Entry(self.root, width=60)
        self.input_field.pack(side=tk.LEFT, padx=10)
        self.input_field.bind("<Return>", self.send_message)
        self.send_button = tk.Button(self.root, text='Send', command=self.send_message)
        self.send_button.pack(side=tk.LEFT)
 
        # Initialize the API key
        self.api_key = ''
 
    def set_api_key(self):
        # 设置API密钥
        self.api_key = tkinter.simpledialog.askstring("Set API Key", "Enter your API Key:")
        if self.api_key:
            self.api_key = self.api_key.strip()
 
    def send_message(self, event=None):
        message = self.input_field.get()
        if not message:
            return
        self.input_field.delete(0, tk.END)
        self.add_chat_message('You: ' + message)
        response = self.get_response(message)
        if response:
            self.add_chat_message('ChatGPT: ' + response)
 
    def add_chat_message(self, message):
        self.chat_window.config(state=tk.NORMAL)
        self.chat_window.insert(tk.END, message + '\n')
        self.chat_window.config(state=tk.DISABLED)
 
    def get_response(self, message):
        if not self.api_key:
            messagebox.showerror('Error', 'API key is not set')
            return
        openai.api_key = self.api_key
        prompt = f"Q: {message}\nA:"
        response = openai.Completion.create(
            engine="davinci",
            prompt=prompt,
            max_tokens=1024,
            n=1,
            stop=None,
            temperature=0.7,
        )
        if response.choices:
            text = response.choices[0].text.strip()
            self.add_chat_message('ChatGPT: ' + text)
            return text
 
 
if __name__ == '__main__':
    root = tk.Tk()
    app = ChatGPTUI(root)
    root.mainloop()