#!/usr/bin/env bash

# Colors
# Used to decorate the replies, queries, error messages etc. used in the script
BLACK='\033[0;30m'
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
PURPLE='\033[0;35m'
CYAN='\033[0;36m'
WHITE='\033[0;37m'
RESET='\033[0m'

# Logo
# Used only for decorative purposes
readonly ASCII=$(
  cat <<EOF
██╗  ██╗ █████╗ ██╗         ██████╗  ██████╗ ██████╗ ██████╗
██║  ██║██╔══██╗██║         ╚════██╗██╔═████╗╚════██╗╚════██╗
███████║███████║██║          █████╔╝██║██╔██║ █████╔╝ █████╔╝
██╔══██║██╔══██║██║         ██╔═══╝ ████╔╝██║██╔═══╝  ╚═══██╗
██║  ██║██║  ██║███████╗    ███████╗╚██████╔╝███████╗██████╔╝
╚═╝  ╚═╝╚═╝  ╚═╝╚══════╝    ╚══════╝ ╚═════╝ ╚══════╝╚═════╝
EOF
)

# Labels
# Decorative line, shown before the answer. It is shown in the following format:
#<NEWLINE>
# ChatGPT ｢<Model Name> | Tokens Used - <Total Tokens Used>｣ (For text based queries)
# ChatGPT ｢Image｣: <Success Message> (For image generation)
# "ChatGPT" shown in bold and cyan. Model Name and Tokens Used are shown in lighter greyed out text.
# CHATGPT - Used for gpt-3.5-turbo model
# CHATGPT_DAVINCI - Used for text-davinci-003 model
# CHATGPT_IMAGE - Used for image generation
readonly CHATGPT="${CYAN}\033[1mChatGPT\033[0m${RESET} \033[2m｢Turbo Model | Tokens Used -\033[0m"
readonly CHATGPT_DAVINCI="${CYAN}\033[1mChatGPT\033[0m${RESET} \033[2m｢Davinci Model | Tokens Used -\033[0m:"
readonly CHATGPT_IMAGE="${CYAN}\033[1mChatGPT\033[0m${RESET} \033[2m｢Image｣\033[0m:"
readonly CHATGPT_GPT4="${CYAN}\033[1mChatGPT\033[0m${RESET} \033[2m｢GPT4 Model | Tokens Used -\033[0m:"

# API Key
# Located at: ~/.chat-gpt-api
# The key is stored unencrypted, in plain text with no spaces etc.
readonly API_KEY=$(<~/.chat-gpt-api)

# Default Parameters
TEMP=${TEMP:-0.8}
MAX_TOKENS=${MAX_TOKENS:-2048}
RESOLUTION=${RESOLUTION:-1024x1024}

# Date & Time
# Used in the history file
readonly DATE_TIME=$(date +"%d-%m-%y at %H:%M")
# Used for Prompts
readonly DATE=$(date +"%d-%m-%Y")

# Pre-defined prompts
readonly PROMPT_CMD="You are an AI assistant. You will generate a SINGLE LINE SHELL COMMAND based on user input. The command should be ready to be executed in the terminal. Show no explanation."

# Message to indicate loading of OpenAI's API response
readonly LOADING="${YELLOW}\033[5mLoading...${RESET}"

# Logo printed when the script is executed
printf "\033c"
echo -e "\n${BLUE}${ASCII}${RESET}\n"

# Checks for ChatGPT Error Message
# response: Response from curl, without jq
# error: To check if response has .error in JSON
# error_message: Specific error message form OpenAI extracted
function connection_error() {
  local response=$1
  local error
  local error_message

  error=$(echo "$response" | jq -r '.error')
  error_message=$(echo "$response" | jq -r '.error.message')

  if [ "$error" != "null" ]; then
    echo -e "\n${RED}Error:${RESET} Your message failed to reach Open AI's API."
    echo -e "${RED}Error Message:${RESET}"
    echo -e "${YELLOW}${error_message}${RESET}\n"
    exit 1
  fi
}

# Parses response from chatGPT using the gpt-3.5-turbo model
# input: User input
# model: ChatGPT Model to be used in the curl request
# response: Response from curl, without jq
# response_parsed: 'response' parsed for message from ChatGPT with text folded to break lines at the max width of the terminal
# response_glow: 'response_parsed' along with 'glow' formatting
# tokens_used: `response` parsed to extract the total number of tokens used for the query
function chat_response() {
  local input=$1
  local model="gpt-3.5-turbo"
  local response
  local response_parsed
  local response_glow
  local tokens_used

  # Loading message while waiting for response
  echo -e "\n$LOADING"

  response=$(curl -sS https://api.openai.com/v1/chat/completions \
    -H "Authorization: Bearer $API_KEY" \
    -H "Content-Type: application/json" \
    -d '{
          "model": "'"$model"'",
          "messages": [
                        {"role": "user", "content": "'"$input"'"}
                      ],
          "max_tokens": '$MAX_TOKENS',
          "temperature": '$TEMP'
          }')

  response_parsed=$(echo "$response" | jq -r '.choices[].message.content' | fold -s -w $COLUMNS)
  tokens_used=$(echo "$response" | jq -r '.usage.total_tokens')

  connection_error "$response"

  # If 'glow' is installed, it is used for formatting
  # Otherwise the output is folded to break lines at the width of the terminal
  if command -v "glow" >/dev/null 2>&1; then
    response_glow=$(echo "$response" | jq -r '.choices[].message.content' | glow)
    # Moves the cursor two line up, and deleted the last two printed line, thus removing the Loading message
    tput cuu1 && tput dl1 && tput cuu1 && tput dl1
    echo -e "${CHATGPT} \033[2m${tokens_used}｣\033[0m:${GREEN}${response_glow}${RESET}"
  else
    tput cuu1 && tput dl1 && tput cuu1 && tput dl1
    echo -e "${CHATGPT} \033[2m${tokens_used}｣\033[0m:${GREEN}${response_parsed}${RESET}"
  fi

  save_history "$input" "$response_parsed" "$model" "$tokens_used"
}

# Parses response from chatGPT using the text-davinci-003 model
# input: User input
# model: ChatGPT Model to be used in the curl request
# response: Response from curl, without jq
# response_parsed: 'response' parsed for message from ChatGPT with text folded to break lines at the max width of the terminal
# response_glow: 'response_parsed' along with 'glow' formatting
# tokens_used: `response` parsed to extract the total number of tokens used for the query
function chat_response_davinci() {
  local input=$1
  local model="text-davinci-003"
  local response
  local response_parsed
  local response_glow
  local tokens_used

  # Loading message while waiting for response
  echo -e "\n$LOADING"

  response=$(curl -sS https://api.openai.com/v1/completions \
    -H 'Content-Type: application/json' \
    -H "Authorization: Bearer $API_KEY" \
    -d '{
    			"model": "'"$model"'",
    			"prompt": "'"$input"'",
    			"max_tokens": '$MAX_TOKENS',
    			"temperature": '$TEMP'
  	}')

  response_parsed=$(echo "$response" | jq -r '.choices[].text' | sed '1,2d' | fold -s -w $COLUMNS)
  tokens_used=$(echo "$response" | jq -r '.usage.total_tokens')

  connection_error "$response"

  # If 'glow' is installed, it is used for formatting
  # Otherwise the output is folded to break lines at the width of the terminal
  if command -v "glow" >/dev/null 2>&1; then
    response_glow=$(echo "$response" | jq -r '.choices[].text' | sed '1,2d' | glow)
    # Moves the cursor two line up, and deleted the last two printed line, thus removing the Loading message
    tput cuu1 && tput dl1 && tput cuu1 && tput dl1
    echo -e "${CHATGPT_DAVINCI} \033[2m${tokens_used}｣\033[0m:${GREEN}${response_glow}${RESET}"
  else
    tput cuu1 && tput dl1 && tput cuu1 && tput dl1
    echo -e "${CHATGPT_DAVINCI} \033[2m${tokens_used}｣\033[0m:${GREEN}${response_parsed}${RESET}"
  fi

  save_history "$input" "$response_parsed" "$model" "$tokens_used"
}

# Generates image
# input: User input
# success_message: String to be displayed along with the URL
# model: ChatGPT Model to be used in the curl request
# response: Response from curl, without jq
# image_url: 'response' parsed for the image URL
# url_short: Shortens `image_url` using tny.im API
function image_generation() {
  local input=$1
  local success_message="URL was successfully generated."
  local model="DALL.E"
  local response

  # Loading message while waiting for response
  echo -e "\n$LOADING"

  response=$(curl -sS https://api.openai.com/v1/images/generations \
    -H 'Content-Type: application/json' \
    -H "Authorization: Bearer $API_KEY" \
    -d '{
          "prompt": "'"$input"'",
          "n": 1,
          "size": "'"$RESOLUTION"'"
    }')

  local image_url
  image_url=$(echo $response | jq -r '.data[0].url')
  local url_short
  url_short=$(wget -qO- -U Mozilla http://tinyurl.com/api-create.php?url=$image_url)

  if [ "$image_url" == null ]; then
    # Moves the cursor two line up, and deleted the last two printed line, thus removing the Loading message
    tput cuu1 && tput dl1 && tput cuu1 && tput dl1
    connection_error "$response"
  else
    tput cuu1 && tput dl1 && tput cuu1 && tput dl1
    echo -e "${CHATGPT_IMAGE} ${GREEN}Image was successfully generated${RESET}"
    echo -e "Link: $url_short"
    save_history "$input" "$success_message" "$model" "NA"
  fi
}

# Parses response from chatGPT using the gpt-4 model
# input: User input
# model: ChatGPT Model to be used in the curl request
# response: Response from curl, without jq
# response_parsed: 'response' parsed for message from ChatGPT with text folded to break lines at the max width of the terminal
# response_glow: 'response_parsed' along with 'glow' formatting
# tokens_used: `response` parsed to extract the total number of tokens used for the query
function chat_response_gpt4() {
  local input=$1
  local model="gpt-4"
  local response
  local response_parsed
  local response_glow
  local tokens_used

  # Loading message while waiting for response
  echo -e "\n$LOADING"

  response=$(curl -sS https://api.openai.com/v1/chat/completions \
    -H "Authorization: Bearer $API_KEY" \
    -H "Content-Type: application/json" \
    -d '{
          "model": "'"$model"'",
          "messages": [
                        {"role": "user", "content": "'"$input"'"}
                      ],
          "max_tokens": '$MAX_TOKENS',
          "temperature": '$TEMP'
          }')

  response_parsed=$(echo "$response" | jq -r '.choices[].text' | sed '1,2d' | fold -s -w $COLUMNS)
  tokens_used=$(echo "$response" | jq -r '.usage.total_tokens')

  connection_error "$response"

  # If 'glow' is installed, it is used for formatting
  # Otherwise the output is folded to break lines at the width of the terminal
  if command -v "glow" >/dev/null 2>&1; then
    response_parsed=$(echo "$response" | jq -r '.choices[].text' | sed '1,2d' | glow)
    # Moves the cursor two line up, and deleted the last two printed line, thus removing the Loading message
    tput cuu1 && tput dl1 && tput cuu1 && tput dl1
    echo -e "${CHATGPT_GPT4} \033[2m${tokens_used}｣\033[0m:${GREEN}${response_glow}${RESET}"
  else
    tput cuu1 && tput dl1 && tput cuu1 && tput dl1
    echo -e "${CHATGPT_GPT4} \033[2m${tokens_used}｣\033[0m:${GREEN}${response_parsed}${RESET}"
  fi

  save_history "$input" "$response_parsed" "$model" "$tokens_used"
}

# Sends pre-defined prompt and parses the single line shell command using the gpt-3.5-turbo Model
# input: User input
# model: ChatGPT Model to be used in the curl request
# response: Response from curl, without jq. Role `system` includes the pre-defined prompt instructing ChatGPT to generate the command based on user input
# response_parsed: 'response' parsed for message from ChatGPT with text folded to break lines at the max width of the terminal
# tokens_used: `response` parsed to extract the total number of tokens used for the query
function command_generation() {
  local input=$1
  local model="gpt-3.5-turbo"
  local response
  local response_parsed
  local tokens_used

  # Loading message while waiting for response
  echo -e "\n$LOADING"

  response=$(curl -sS https://api.openai.com/v1/chat/completions \
    -H "Authorization: Bearer $API_KEY" \
    -H "Content-Type: application/json" \
    -d '{
          "model": "'"$model"'",
          "messages": [
                        {"role": "system", "content": "'"$PROMPT_CMD"'"},
                        {"role": "user", "content": "'"$input"'"}
                      ],
          "max_tokens": '$MAX_TOKENS',
          "temperature": '$TEMP'
          }')

  response_parsed=$(echo "$response" | jq -r '.choices[].message.content' | fold -s -w $COLUMNS)
  tokens_used=$(echo "$response" | jq -r '.usage.total_tokens')

  connection_error "$response"
    # Moves the cursor two line up, and deleted the last two printed line, thus removing the Loading message
  tput cuu1 && tput dl1 && tput cuu1 && tput dl1
  echo -e "${CHATGPT} \033[2m${tokens_used}｣\033[0m:\n${GREEN}${response_parsed}${RESET}"
  save_history "$input" "$response_parsed" "$model" "$tokens_used"
}

# Takes user input to set the tone and length for specific content generation
# input: Takes user input regarding the specific content to be generated, email for example
#        Based on the user input (email, article etc.) the generated information from this function is forwarded to the
#        relevant function.
# length: Length of the content to be generated
# tone: Tone of the content to be generated
# prompt_init: Initial prompt generated including today's date along with the length and tone (if entered by the user)
function prompt_length_tone() {
  local input=$1
  local length
  local tone
  local prompt_init
  local prompt_final

  prompt_init="Today is ${DATE}."

  echo -e "\nWhat length should the content be?"
  echo -e "Press ${YELLOW}1${RESET} for short."
  echo -e "Press ${YELLOW}2${RESET} for medium."
  echo -e "Press ${YELLOW}3${RESET} for long."
  echo -e "\nEnter any value if you would like to leave the length unspecified."

  # Reads the input length in raw mode + readline mode
  read -re length

  if [ "$length" -eq 1 ]; then
    length="short"
  elif [ "$length" -eq 2 ]; then
    length="medium"
  elif [ "$length" -eq 3 ]; then
    length="long"
  else
    length=""
  fi

  echo -e "\nWhat tone should the content have?"
  echo -e "Enter ${YELLOW}1${RESET} for casual."
  echo -e "Enter ${YELLOW}2${RESET} for funny."
  echo -e "Enter ${YELLOW}3${RESET} for informational."
  echo -e "Enter ${YELLOW}4${RESET} for professional."
  echo -e "\nEnter any other value if you would like to leave the tone unspecified."

  # Reads the input for tone in raw mode + readline mode
  read -re tone

  if [ "$tone" -eq 1 ]; then
    tone="casual"
  elif [ "$tone" -eq 2 ]; then
    tone="funny"
  elif [ "$tone" -eq 3 ]; then
    tone="informational"
  elif [ "$tone" -eq 4 ]; then
    tone="professional"
  else
    tone=""
  fi

  if [[ -z "$length" && -z "$tone" ]]; then
    prompt_final="$prompt_init"
  elif [[ -n "$length" && -n "$tone" ]]; then
    prompt_final="$prompt_init The length should be $length. The tone of the email should be $tone."
  elif [ -n "$length" ]; then
    prompt_final="$prompt_init The length should be $length."
  elif [ -n "$tone" ]; then
    prompt_final="$prompt_init The tone of the email should be $tone"
  fi

  if [ "$input" == "-email" ] || [ "$input" == "--e" ]; then
    email_generation "$prompt_final"
  elif [ "$input" == "-article" ] || [ "$input" == "--a" ]; then
    article_generation "$prompt_final"
  fi

}

# Prepares a prompt for email drafting and forwards it to gpt-3.5-turbo model
# prompt_init: Sent from prompt_length_tone() function and includes the current date, length and tone for the query
# prompt_info: Directs ChatGPT to generate email with the user query
# input: User input of the actual query
# final_query: Concatenates prompt_init and user input. This is sent to chat_response() function
function email_generation() {
  local prompt_init=$1
  local prompt_info
  local final_query

  prompt_info="Draft an email. $prompt_init"

  echo -e "\n${WHITE}Write a prompt for your email...${RESET}"
  echo -e "${WHITE}Include any and all details you would like for the email to have.${RESET}\n"

  # Reads the input in raw mode + readline mode
  read -re input

  # Tone and length (as selected by the user) along with the query input concatenated into "final_query"
  final_query="$prompt_info $input"

  echo -e "\n${PURPLE}Generating response for:${RESET}"
  echo -e "${BLUE}$final_query${RESET}"

  # "final_query" sent to gpt-3.5-turbo model
  chat_response "$final_query"
}

# Prepares a prompt for article drafting and forwards it to gpt-3.5-turbo model
# prompt_init: Sent from prompt_length_tone() function and includes the current date, length and tone for the query
# prompt_info: Directs ChatGPT to generate email with the user query
# input: User input of the actual query
# final_query: Concatenates prompt_init and user input. This is sent to chat_response() function
function article_generation() {
  local prompt_init=$1
  local prompt_info
  local final_query

  prompt_info="Draft an article as concisely as possible. $prompt_init"

  echo -e "\n${WHITE}Write a prompt for your article...${RESET}"
  echo -e "${WHITE}Include any and all details you would like for the article to have, for example${RESET}"
  echo -e "${WHITE}intended platform, number of paragraphs, target audience and so on.${RESET}\n"

  # Reads the input in raw mode + readline mode
  read -re input

  # Tone and length (as selected by the user) along with the query input concatenated into "final_query"
  final_query="$prompt_info $input"

  echo -e "\n${PURPLE}Generating response for:${RESET}"
  echo -e "${BLUE}$final_query${RESET}"

  # "final_query" sent to gpt-3.5-turbo model
  chat_response "$final_query"
}

# Saves responses in a history file with timestamp
# Located at: ~/.hal2023_history.txt
# History file isn't deleted while uninstalling
# prompt: User input
# reply: Parsed reply from ChatGPT to the user input
# model: Model used for the query
# tokens_used: Tokens used for the text based query. Image generation does not display the total number of tokens used
function save_history() {
  local prompt=$1
  local reply=$2
  local model=$3
  local tokens_used=$4

  # For each individual query, the above variables are stored in the following format:
  # <NEWLINE>
  # +++++++++++++++++ (used for decorative purposes)
  # <NEWLINE>
  # <NEWLINE>
  # dd-mm-yy at HH:MM
  # User: <User query>
  # <NEWLINE>
  # ChatGPT [<Model>] [Tokens Used - <Tokens>] (For image generation tokens used is displayed as NA)
  # <NEWLINE>
  # <ChatGPT Reply> (For image generation "URL was successfully generated" is displayed"
  echo -e "\n+++++++++++++++++\n\n$DATE_TIME\n\nUser: $prompt \n\nChatGPT [$model] [Tokens Used - $tokens_used]\n$reply" >>~/.hal2023_history.txt
}

# Checks for a History File during script run
# If History file is not found, it is created
if [ ! -f ~/.hal2023_history.txt ]; then
  touch ~/.hal2023_history.txt
fi

# Endless `while` loop that runs until the user types in one of the exit keywords
# By default `gpt-3.5-turbo` model is used to reply to queries
# "-davinci" keyword before a prompt points the query to the `text-davinci-003` model via the chat_response_davinci() function
# "-image" keyword before a prompt points the query to the `DALL.E` model for image generation via the image_generation() function
# "-email" or "--e" keyword points the user to set the length, tone and then finally prompt input for email generation using the 'gpt-3.5-turbo'
# "-article" or "--a" keyword points the user to set the length, tone and then finally prompt input for article generation using the 'gpt-3.5-turbo'
while true; do
  echo -e "\n\033[2mType exit, quit or :q to exit.\033[0m"
  echo -e "\033[3mWrite a query and press enter twice\033[0m\n"

  # Reads the input in raw mode + readline mode
  # Press ENTER twice (empty newline) to accept input
  # Any trailing whitespaces and newlines are
  unset REPLY input
  while read -re; do
    [[ -n "$REPLY" ]] || break
    input="$input${input:+\\n}$REPLY"
  done; input="${input:-$REPLY}"

  if [ "$input" == "exit" ] || [ "$input" == "quit" ] || [ "$input" == ":q" ]; then
    echo -e "\nSuccessfully quit HAL2023\n"
    exit 0
  # Input sent to text-davinci-003 model
  elif [[ "$input" == "-davinci"* ]]; then
    davinciQuery=$(echo "$input" | cut -c 9-)
    chat_response_davinci "$davinciQuery"
  # Input sent to image_generation() function
  elif [[ "$input" == "-image"* ]]; then
    image_generation "${input#*-image }"
  # Input sent to gpt-4 model
  elif [[ "$input" == "-gpt4"* ]]; then
    chat_response_gpt4 "${input#*-gpt4 }"
  # Rerouted to prompt_length_tone() function and then subsequently to email_generation() function
  elif [ "$input" == "-email" ] || [ "$input" == "--e" ]; then
    prompt_length_tone "$input"
  # Rerouted to prompt_length_tone() function and then subsequently to article_generation() function
  elif [ "$input" == "-article" ] || [ "$input" == "--a" ]; then
    prompt_length_tone "$input"
  # Input sent to command_generation() function and passes along with pre-defined prompt to generate a single line bash command
  elif [[ "$input" == "-cmd"* ]]; then
    command_generation "${input#*-cmd }"
  else
    chat_response "$input"
  fi
done
