#!/usr/bin/env python3

import json, os, re, sys
from openai import OpenAI

try:
	from dotenv import load_dotenv
	load_dotenv()
except ImportError:
	pass


api_key = os.getenv('OPENAI_API_KEY')
model = os.getenv('OPENAI_MODEL')
default_prompt = os.getenv('OPENAI_DEFAULT_PROMPT')
log_file = os.getenv('OPENAI_LOGFILE')
temperature = os.getenv('OPENAI_TEMPERATURE')
max_tokens = os.getenv('OPENAI_MAX_TOKENS')


if not model:
	model = 'gpt-3.5-turbo'
if not default_prompt:
	default_prompt = ''
if not temperature:
	temperature = 0.0
else:
	temperature = float(temperature)
if not max_tokens:
	max_tokens = 2048
else:
	max_tokens = int(max_tokens)

client = OpenAI(api_key=api_key)

def write_to_log(display: str):
	with open(log_file, 'a') as f:
		f.write(f"{display}\n")

def get_answer(prompt: str, response=None, choices=None, result=None) -> str:
	response = client.chat.completions.create(
		model=model,
		messages=[
            		{"role": "system", "content": default_prompt},
            		{"role": "user", "content": prompt}
        	],
		temperature=temperature,
		max_tokens=max_tokens
		)
	choices = response.choices
	if choices:
		res = choices[0].message.content
	if res:
		return re.sub(r'^\s+|\s+$', '', res)  # Trim leading/trailing whitespace
	return None


if __name__ == '__main__':
	if not sys.argv[1:]:
		print("ERROR: gpt must run with [quoted] string of text as prompt.")
		sys.exit()
	real_prompt = ''.join(sys.argv[1:])
	prompt = f"{default_prompt}{real_prompt}"
	answer = get_answer(prompt)
	if not answer:
		print("Error retrieving answer.")
		sys.exit()
	display = (
			"Prompt:\n"
			f"{real_prompt}\n\n"
			"Answer:\n"
			f"{answer}\n"
			"----------------------------------------------------------------------"
		)
	print(display)
	if log_file:
		write_to_log(display)
