#!/bin/bash
# ================================================================
# Nuvoton ChatGPT demo
# Run ChatGPT demo using text on Nuvoton platform
# 
# SPDX-License-Identifier: Apache-2.0
# Copyright (C) 2024 Nuvoton Technology Corp. All rights reserved.
# ================================================================


CHAT_INIT_PROMPT="You are ChatGPT, a Large Language Model trained by OpenAI. You will be answering questions from users. You answer as concisely as possible for each response (e.g. don’t be verbose). If you are generating a list, do not have too many items. Keep the number of items short. Before each user prompt you will be given the chat history in Q&A form. Output your answer directly, with no labels in front. Do not start your answers with A or Anwser. You were trained on data up until 2023. Today's date is $(date +%m/%d/%Y)"

SYSTEM_PROMPT="You are ChatGPT, a large language model trained by OpenAI. Answer as concisely as possible. Current date: $(date +%m/%d/%Y). Knowledge cutoff: 10/1/2023."

CHATGPT_CYAN_LABEL="\033[36mchatgpt \033[0m"
PROCESSING_LABEL="\033[95mProcessing... \033[0m\033[0K\r"
OVERWRITE_PROCESSING_LINE="             \033[0K\r"

if [[ -z "$OPENAI_KEY" ]]; then
	echo "You need to set your OPENAI_KEY to use this script"
	echo "You can set it temporarily by running this on your terminal: export OPENAI_KEY=YOUR_KEY_HERE"
	exit 1
fi

usage() {
	cat <<EOF
A simple, lightweight shell script to use OpenAI's Language Models and DALL-E from the terminal.

Commands:
  image: - To generate an image and store it named "image.png" in local, you also can open the image in your browser.

Options:
  -p, --prompt               Provide prompt instead of starting chat

  -l, --list                 List available openAI models

  -m, --model                Model to use

  -t, --temperature          Temperature

  -s, --size                 Image size. (The sizes that are accepted by the OpenAI API is
                             dall-e-3 model: 1024x1024, 1792x1024 or 1024x1792,
                             dall-e-2 model: 256x256, 512x512 or 1024x1024)

EOF
}

# error handling function
# $1 should be the response body
handle_error() {
	if echo "$1" | jq -e '.error' >/dev/null; then
		echo -e "Your request to Open AI API failed: \033[0;31m$(echo "$1" | jq -r '.error.type')\033[0m"
		echo "$1" | jq -r '.error.message'
		exit 1
	fi
}

# request to openAI API models endpoint. Returns a list of models
# takes no input parameters
list_models() {
	models_response=$(curl https://api.openai.com/v1/models \
		-sS \
		-H "Authorization: Bearer $OPENAI_KEY")
	handle_error "$models_response"
	models_data=$(echo $models_response | jq -r -C '.data[] | {id, owned_by, created}')
	echo -e "$OVERWRITE_PROCESSING_LINE"
	echo -e "${CHATGPT_CYAN_LABEL}This is a list of models currently available at OpenAI API:\n ${models_data}"
}

# escapes special characters to ensure the string can be safely used
# in a JSON-compatible format
# $1 should be the message(s) to do escape
escape() {
	echo "$1" | jq -Rrs 'tojson[1:-1]'
}

# build user chat message function for /chat/completions (gpt models)
# builds chat message before request,
# $1 should be the escaped request prompt,
# it extends $chat_message
build_user_chat_message() {
	local escaped_request_prompt="$1"
	if [ -z "$chat_message" ]; then
		chat_message="{\"role\": \"user\", \"content\": \"$escaped_request_prompt\"}"
	else
		chat_message="$chat_message, {\"role\": \"user\", \"content\": \"$escaped_request_prompt\"}"
	fi
}

# adds the assistant response to the message in (chatml) format
# for /chat/completions (gpt models)
# keeps messages length under max token limit,
# $1 should be the escaped response data
# it extends and potentially shrinks $chat_message
add_assistant_response_to_chat_message() {
	local escaped_response_data="$1"
	# add response to chat context as answer
	chat_message="$chat_message, {\"role\": \"assistant\", \"content\": \"$escaped_response_data\"}"

	# check prompt length, 1 word =~ 1.3 tokens
	# reserving 100 tokens for next user prompt
	while (($(echo "$chat_message" | wc -c) * 13 / 10 > (MAX_TOKENS - 100))); do
		# transform to json array to parse with jq
		local chat_message_json="[ $chat_message ]"
		# remove first/oldest QnA from prompt
		local chat_message_removed=$(echo "$chat_message_json" | jq -c '.[2:]')
		# remove the json array format
		chat_message=${chat_message_removed:1:-1}
	done
}

# request to OpenAPI API chat completion endpoint function
# $1 should be the message(s) formatted with role and content
request_to_chat() {
	local message="$1"
	escaped_system_prompt=$(escape "$SYSTEM_PROMPT")

	curl https://api.openai.com/v1/chat/completions \
		-sS \
		-H 'Content-Type: application/json' \
		-H "Authorization: Bearer $OPENAI_KEY" \
		-d '{
			"model": "'"$MODEL"'",
			"messages": [
				{"role": "system", "content": "'"$escaped_system_prompt"'"},
				'"$message"'
			],
			"max_completion_tokens": '$MAX_COMPLETION_TOKENS',
			"temperature": '$TEMPERATURE'
		}'
}

# request to OpenAI API image generations endpoint function
# $1 should be the prompt
request_to_image() {
	local prompt="$1"
	image_response=$(curl https://api.openai.com/v1/images/generations \
		-sS \
		-H 'Content-Type: application/json' \
		-H "Authorization: Bearer $OPENAI_KEY" \
		-d '{
			"model": "dall-e-3",
			"prompt": "'"${prompt#*image:}"'",
			"n": 1,
			"size": "'"$SIZE"'"
		}')
}

# function to wrap text considering different languages
# $1 should be the response body
text_display() {
echo "$1" | gawk -v max_width=$MAX_WIDTH '
{
	# initialize variables
	line = ""  # current line content
	len = 0	# current line length

	# split the line into words by spaces
	n = split($0, words, " ")

	# loop through each word
	for (i = 1; i <= n; i++) {
		word = words[i]  # extract the current word
		word_len = 0	 # length of the current word (consider ASCII and non-ASCII characters)

		# calculate the length of the word (English words and non-ASCII characters)
		for (j = 1; j <= length(word); j++) {
			char = substr(word, j, 1)

			# check if the character is an ASCII character (1 space)
			if (char ~ /[ -~]/) {
				word_len += 1
			} else {
				# non-ASCII characters (Chinese, Japanese) take 2 spaces
				word_len += 2
			}
		}

		# handle English words without splitting them
		if (word_len == length(word)) {
			# if adding the entire word exceeds max_width spaces, print the current line and start a new one
			if (len + word_len + (len > 0 ? 1 : 0) > max_width) {
				print line
				line = word
				len = word_len
			} else {
				# add the word to the current line
				if (len > 0) {
					line = line " " word  # add space between words
					len += 1
				} else {
					line = word
				}
				len += word_len
			}
		} else {
			# process Chinese/Japanese characters individually
			for (j = 1; j <= length(word); j++) {
				char = substr(word, j, 1)

				# determine the length of the character
				if (char ~ /[ -~]/) {
					char_len = 1
				} else {
					char_len = 2
				}

				# if adding this character exceeds max_width spaces, print the current line and start a new one
				if (len + char_len > max_width) {
					print line
					line = char
					len = char_len
				} else {
					# add the character to the current line
					line = line char
					len += char_len
				}
			}
		}
	}

	# print any remaining characters in the line
	if (len > 0) {
		print line
	}
}
'
}

# parse command line arguments
while [[ "$#" -gt 0 ]]; do
	case $1 in
	-p | --prompt)
		prompt="$2"
		shift
		shift
		;;
	-l | --list)
		list_models
		exit 0
		;;
	-m | --model)
		MODEL="$2"
		shift
		shift
		;;
	-t | --temperature)
		TEMPERATURE="$2"
		shift
		shift
		;;
	-s | --size)
		SIZE="$2"
		shift
		shift
		;;
	-h | --help)
		usage
		exit 0
		;;
	*)
		echo "Unknown parameter: $1"
		exit 1
		;;
	esac
done

# set defaults
MODEL=${MODEL:-gpt-4o}
MAX_TOKENS=${MAX_TOKENS:-4096}
MAX_COMPLETION_TOKENS=${MAX_COMPLETION_TOKENS:-4096}
TEMPERATURE=${TEMPERATURE:-0.7}
MAX_WIDTH=${MAX_WIDTH:-80}
SIZE=${SIZE:-1024x1024}

# create history file
if [ ! -f ~/.chatgpt_history ]; then
	touch ~/.chatgpt_history
	chmod 600 ~/.chatgpt_history
fi

# check input source and determine run mode
running=true
# show the time information in processing
time_info=true

# prompt from argument, run on pipe mode (run once, no chat)
if [ -n "$prompt" ]; then
	pipe_mode_prompt=${prompt}
# if input file_descriptor is a terminal, run on chat mode
elif [ -t 0 ]; then
	echo -e "Welcome to chatgpt. You can quit with '\033[36mexit\033[0m' or '\033[36mq\033[0m'."
# prompt from pipe or redirected stdin, run on pipe mode
else
	pipe_mode_prompt+=$(cat -)
fi

while $running; do

	if [ -z "$pipe_mode_prompt" ]; then
		echo -e "\nEnter a prompt:"
		read -e prompt
		if [[ ! $prompt =~ ^(exit|q)$ ]]; then
			if [ "$time_info" = true ]; then
				echo -ne "\033[33m["`date "+%H:%M:%S"`"] \033[0m"
			fi
			echo -e $PROCESSING_LABEL
		fi
	else
		# set vars for pipe mode
		prompt=${pipe_mode_prompt}
		running=false
		CHATGPT_CYAN_LABEL=""
	fi

	if [[ $prompt =~ ^(exit|q)$ ]]; then
		running=false
	elif [[ "$prompt" =~ ^image: ]]; then
		request_to_image "$prompt"
		handle_error "$image_response"
		image_url=$(echo "$image_response" | jq -r '.data[0].url')
		if [ "$time_info" = true ]; then
			echo -ne "\n\033[33m["`date "+%H:%M:%S"`"] \033[0m"
		fi
		echo -e "${CHATGPT_CYAN_LABEL}Your image was created. \n\nLink: ${image_url}\n"

		if [ -f image.png ]; then
			rm image.png
		fi
		curl -sS $image_url -o image.png
		if [ "$time_info" = true ]; then
			echo -ne "\033[33m["`date "+%H:%M:%S"`"] \033[0m"
		fi
		echo -e "${CHATGPT_CYAN_LABEL}Your image was downloaded."
	else
		# escape quotation marks, new lines, backslashes...
		request_prompt=$(escape "$prompt")

		build_user_chat_message "$request_prompt"
		response=$(request_to_chat "$chat_message")
		handle_error "$response"
		response_data=$(echo "$response" | jq -r '.choices[].message.content')

		if [ "$time_info" = true ]; then
			echo -ne "\n\033[33m["`date "+%H:%M:%S"`"] \033[0m"
		fi
		echo -e "${CHATGPT_CYAN_LABEL}"
		text_display "${response_data}"
		echo -e "${CHATGPT_CYAN_LABEL}${response_data}" > response.txt
		add_assistant_response_to_chat_message "$(escape "$response_data")"

		timestamp=$(date +"%Y-%m-%d %H:%M")
		echo -e "$timestamp $prompt \n$response_data \n" >>~/.chatgpt_history
	fi
done
