#!/bin/bash

# Examples of reading the arguments
readonly arg_count=$#
readonly first_arg=$1
readonly all_args=$*

# Construct and make an OpenAI query
openai_query () {

	# Check if a query has been provided, otherwise use a default
	local prompt="haiku about c++"
	[[ -z $* ]] || prompt=$*

	# Models
	# The first is the most accurate if a little slow
	# But it does fail a lot, perhaps because it gets hammered as the default?

	# text-davinci-003
	# text-davinci-002
	# text-davinci-001
	# text-curie-003

	# Create request JSON
	request=$(cat <<!
{
	"model": "gpt-3.5-turbo-instruct",
	"prompt": "$prompt",
	"temperature": 0.7,
	"max_tokens": 1024,
	"top_p": 1,
	"frequency_penalty": 0.0,
	"presence_penalty": 0.0
}
!
)

	# Issue the request
	local max_curl_time=20
	curl --silent --max-time $max_curl_time https://api.openai.com/v1/completions \
	  -H "Content-Type: application/json" \
	  -H "Authorization: Bearer $OPENAI_API_KEY" \
	  -d "$request"
	
	# Use the OpenAI JSON format for curl errors
	local curl_response=$?
	[[ $curl_response -gt 0 ]] && cat <<!
{
  "error": {
    "message": "curl error",
    "type": "network_error",
    "param": $max_curl_time,
    "code": $curl_response
  }
}
!
}

# Check the response for various errors
validate_response () {

	local response="$*"

	# Extract the response text if it was successful
	# Not we're swallowing two newlines in the response
	if [[ "$response" =~ \"text\":\ \"\\n\\n(.*)\", ]]; then

		echo "${BASH_REMATCH[1]}"
		return
	fi

	# Otherwise dump the error JSON
	echo "$response"
}

# Make query and retry if there's a problem
openai_query_with_retry () {

	local prompt=$*
	local response="not run"

	for i in {1..10}; do

		# Make query and check it
		response1=$(openai_query "$prompt")
		response2=$(validate_response "$response1")

		# Retry if there's an error
		# Debug goes to stderr
		if [[ "$response2" =~ \"error\": ]]; then
			echo "$response2"
			echo "Retry $i: $prompt"
			continue
		fi >&2

		# Otherwise we're good, but trim the first two newlines of the response
		echo -e "$response2"
		return
	done

	# Report the error response if we've got here after retrying
	# This error JSON will end up in the input file instead of the response
	echo "COMPUTER_SAYS_NO: $response2"
}

# Process a simple command line query
if [[ $arg_count -ge 2 ]]; then
	
	# Send all command line params as a prompt
	prompt=$all_args
	openai_query_with_retry "$prompt"

	exit 0
fi

# Create a tmp file in case we're overwriting the input file
readonly temp_file=$(mktemp)

# Process a named file or stdin
if [[ $arg_count -eq 0 ]] || [[ $arg_count -eq 1 ]]; then

	# Preserve leading spaces
	# Otherwise we lose the indents in code excerpts
	IFS=''

	# By magic this works for both stdin and a file
	# If first arg is undefined, cat just passes on stdin
	cat $first_arg | while read -r line; do

		# Check if the line starts and ends with a brace
		if [[ $line =~ ^\{(.*)\}$ ]]; then

			# Use the bit between the braces as the prompt
			prompt=${BASH_REMATCH[1]}
			openai_query_with_retry "$prompt"
		else
			echo "$line"
		fi
	done

fi | tee "$temp_file"

if [[ $arg_count -eq 1 ]]; then
	cp "$temp_file" "$first_arg"
	echo cp "$temp_file" "$first_arg" >&2
fi

# Tidy up after ourselves
rm "$temp_file"
echo rm "$temp_file" >&2