#!/usr/bin/env python3
import sys
import os
import subprocess
import json
import time
if len(sys.argv) < 3: print("Arguments: <text editor> <topic>"); sys.exit(1)
prettyDistro = "simpleGPT"
distro = "me.fier." +prettyDistro
if "APPDATA" in os.environ: userDir = os.environ["APPDATA"]
if "HOME" in os.environ: userDir = os.path.join(os.environ["HOME"],".config")
userDir = os.path.join(userDir,distro)
os.makedirs(os.path.join(userDir,"conversations"),exist_ok = True)

api_key = False
if "OPENAI_API_KEY" in os.environ:
	api_key = os.environ["OPENAI_API_KEY"]
elif os.path.isfile(os.path.join(userDir,"api_key.txt")):
	api_key = open(os.path.join(userDir,"api_key.txt"),"r").read().strip(" \t\r\n")

if api_key == False:
	print("Couldn't find API key, please set OPENAI_API_KEY or create " +os.path.join(userDir,"api_key.txt"))
	os.exit(1)

rolePrefix = "* "
roleSuffix = " says:"
roles = ["system","user","assistant"]

def convertListToHumanReadable(lst):
	txt = ""
	for message in lst:
		txt += rolePrefix +message["role"]+ roleSuffix + "\n" +message["content"].strip("\n")+ "\n\n" 
	return txt

def convertHumanReadableToList(txt):
	lst = []
	txt = txt.split("\n")
	message = False
	for line in txt:
		isrole = False
		for role in roles:
			if line == rolePrefix + role + roleSuffix:
				if message != False:
					message["content"] = message["content"].strip("\n")
					lst.append(message)
				
				message = {"role": role, "content": ""}
				isrole = True
				break
		
		if isrole: continue
		message["content"] += line + "\n"
	lst.append(message)
	return lst

import urllib.parse
import urllib.request
def chatCompletion(lst):
	lstOut = {}
	lstOut["model"] = "gpt-3.5-turbo"
	lstOut["messages"] = lst
	data = json.dumps(lstOut).encode("utf-8")
	dataLength = len(data)
	request = urllib.request.Request("https://api.openai.com/v1/chat/completions",data=data)
	request.add_header("Content-Type","application/json")
	request.add_header("Authorization","Bearer " + api_key)
	response = urllib.request.urlopen(request)
	rtn = json.loads(response.read().decode("utf-8"))
	response.close()
	return rtn

userFile = os.path.join(userDir,"conversations",sys.argv[2] + ".txt")
userFileTmp = userFile + ".tmp"
if not os.path.isfile(userFile):
	lst = []
	lst.append({"role": "system", "content": "You are a smart and helpful chat assistant."})
	lst.append({"role": "user", "content": ""})
	open(userFile,"w",encoding="utf-8").write(convertListToHumanReadable(lst).strip(" \t\r\n") + "\n")

while True:
	subprocess.Popen([sys.argv[1],userFile]).wait()
	lst = convertHumanReadableToList(open(userFile,"r",encoding="utf8").read())
	lst[-1]["content"] = lst[-1]["content"].strip(" \t\r\n")
	if not (lst[-1]["role"] in ["user","system"]) or lst[-1]["content"] == "": sys.exit(0)
	print("Awaiting response ...")
	response = chatCompletion(lst)
	lst.append({"role": "assistant", "content": response["choices"][0]["message"]["content"]})
	lst.append({"role": "user", "content": ""})
	open(userFileTmp,"w",encoding="utf-8").write(convertListToHumanReadable(lst).strip(" \t\r\n") + "\n")
	os.replace(userFileTmp,userFile)