chat

#!/usr/bin/env bash

set -eu -o pipefail

# ------------------------------------------------------------------------------
# Params
# ------------------------------------------------------------------------------
PROGRAM="${0##*/}"
STORE="${BASHGPT_CHAT_HOME:-$HOME/.bashgpt/chat}"
CHAT_HISTORY_FILE="$STORE/chat-history"
LAST_ASSISTANT_FILE="$STORE/last_assistant.conf"
MAX_FILE_SIZE=30000
MAX_MSG_SIZE=15000

# ------------------------------------------------------------------------------
# Globals
# ------------------------------------------------------------------------------
ASSISTANT_NAME=""
ASSISTANT_ID=""
THREAD_ID=""
THREAD_LABEL=""
MESSAGE_ID=""

CONVERSATION_TMP=$(mktemp)
trap 'rm -f $CONVERSATION_TMP' TERM

# ------------------------------------------------------------------------------
# Functions
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# Verifies that the tools used by this script are present and that the required
# environment variables are set. Creates our storage directory if it doesn't
# exist. Finally, it restores the last selected assistant, if any.
#
# Note that `echo` is used here instead of our output wrappers because we have
# not yet ascertained that `gum` is available.
# ------------------------------------------------------------------------------
validate-env() {
  local tools=("curl" "jq" "gum" "tac" "openai")
  local cmd

  for cmd in "${tools[@]}"; do
    if ! command -v "$cmd" &> /dev/null; then
      echo "$cmd is not installed" >&2

      if [ "$cmd" == "openai" ]; then
        echo " - $cmd is a tool provided by this project." >&2
        echo ' - You can make it available by putting it on your PATH' >&2
        # shellcheck disable=SC2016
        echo '      export PATH="$PATH:/path/to/bash-gpt"' >&2
      elif [ "$cmd" == "tac" ]; then
        echo " - $cmd is a GNU coreutils tool." >&2
        echo " - You can install it with 'brew install coreutils'" >&2
      fi

      exit 1
    fi
  done

  if [ -z "${OPENAI_API_KEY:-}" ]; then
    die "OPENAI_API_KEY is not set in the environment"
  fi

  if gum --version | grep -q -v 'v0.12.0'; then
    warn "Version warning: gum v0.12.0 is suggested:"
    warn "  - v0.13 broke 'escape' in 'gum write'"
    warn "  - v0.14 broke *everything* in 'gum write'"
  fi

  mkdir -p "$STORE"

  touch "$CHAT_HISTORY_FILE"
  touch "$LAST_ASSISTANT_FILE"

  restore-selected-assistant

  return 0
}

# ------------------------------------------------------------------------------
# Output wrappers that print messages in pretty colors. All of them send their
# output to STDERR, not that this was ever intended for use in a pipeline.
# ------------------------------------------------------------------------------
info() { gum style --foreground="#66FF66" "$PROGRAM: $1" >&2; }
warn() { gum style --foreground="#FFFF00" "$PROGRAM: ${1:-something went wrong}" >&2; }
die() { gum style --foreground="#FF0000" "$PROGRAM: ${1:-something went wrong}" exit 1 >&2; }

# ------------------------------------------------------------------------------
# API calls
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# Gets a regular completion using the old API. This is only really used for
# having gpt-3.5-turbo-16k summarize a conversation.
# ------------------------------------------------------------------------------
get-completion() {
  local model="$1"
  local system="$2"
  local prompt="$3"
  local result

  if ! result=$(openai get-completion --model "$model" --system "$system" --prompt "$prompt"); then
    warn "error getting completion"
    die "$result"
  fi

  jq -r '.choices[0].message.content' <<< "$result"
}

# ------------------------------------------------------------------------------
# Retrieves the list of assistants from the API and prints them to STDOUT in the
# format "id:name".
# ------------------------------------------------------------------------------
list-assistants() {
  local result

  if ! result=$(openai list-assistants); then
    warn "error listing assistants"
    die "$result"
  fi

  while IFS=$'\t' read -r id name; do
    echo "$id:$name"
  done < <(jq -r '.data[] | "\(.id)\t\(.name // "null")"' <<< "$result")
}

# ------------------------------------------------------------------------------
# Creates a new assistant using the API and prints the ID of the new assistant.
# ------------------------------------------------------------------------------
create-assistant() {
  local name="$1"
  local instructions="$2"
  local description="$3"
  local model="$4"

  # Now include the tool arguments in the openai invocation
  if ! result=$(
    openai create-assistant \
      --name "$name" \
      --instructions "$instructions" \
      --description "$description" \
      --model "$model" \
      2>&1
  ); then
    die "$result"
  fi

  jq -r '.id' <<< "$result"
}

# ------------------------------------------------------------------------------
# Modifies an existing assistant using the API and prints the ID of the modified
# assistant.
# ------------------------------------------------------------------------------
modify-assistant() {
  local name="$1"
  local instructions="$2"
  local description="$3"
  local model="$4"
  local result

  # Now include the tool arguments in the openai invocation
  if ! result=$(
    openai modify-assistant \
      --assistant "$ASSISTANT_ID" \
      --name "$name" \
      --instructions "$instructions" \
      --description "$description" \
      --model "$model" \
      2>&1
  ); then
    die "$result"
  fi

  jq -r '.id' <<< "$result"
}

# ------------------------------------------------------------------------------
# Deletes an existing assistant using the API.
# ------------------------------------------------------------------------------
delete-assistant() {
  local result

  if ! result=$(openai delete-assistant --assistant "$ASSISTANT_ID"); then
    warn "error deleting assistant"
    die "$result"
  fi

  ASSISTANT_ID=""
  ASSISTANT_NAME=""

  return 0
}

# ------------------------------------------------------------------------------
# Creates a new thread and prints out its id.
# ------------------------------------------------------------------------------
start-thread() {
  local result
  local id

  if ! result=$(openai start-thread); then
    warn "error creating a new thread"
    die "$result"
  fi

  jq -r '.id' <<< "$result"
}

# ------------------------------------------------------------------------------
# Appends a message to the current thread.
# ------------------------------------------------------------------------------
send-msg() {
  local msg="$1"
  local result

  if ! result=$(openai add-message --thread "$THREAD_ID" --content "$msg"); then
    warn "error sending message"
    die "$result"
  fi

  print-user-msg "$msg"
}

# ------------------------------------------------------------------------------
# Retrieves GPT's response to the current thread, streaming the output to
# STDOUT.
# ------------------------------------------------------------------------------
get-response() {
  print-gpt-msg
  echo
  openai run-thread --assistant "$ASSISTANT_ID" --thread "$THREAD_ID" --stream
  echo
}

# ------------------------------------------------------------------------------
# Polls the API for the current status of a thread run.
# ------------------------------------------------------------------------------
# Called in bg proc using the spinner
# shellcheck disable=2317
get-run-status() {
  local run_id="$1"
  local result

  if ! result=$(openai get-run --thread "$THREAD_ID" --run "$run_id"); then
    warn "error getting run"
    die "$result"
  fi

  jq -r '.status' <<< "$result"
}

# ------------------------------------------------------------------------------
# Retrieves the messages for the current thread from the API, invoking the
# user-specified callback function for each message received. The callback
# function is passed the message role, and text of the message.
#
# Accepts two optional parameters to limit the number of messages received,
# `limit` (pos int) and `get_all` (bool). These are really only used to
# retrieve the first few messages when generating a summary line for the
# thread. If `get_all` is false, it will restore the original state of
# MESSAGE_ID (since that param is used only for generated summaries, and we
# don't want that to affect state for the chat loop).
# ------------------------------------------------------------------------------
list-messages() {
  local callback=$1
  local messages
  local role
  local text
  local has_more=true
  local lines

  while [ "$has_more" == true ]; do
    messages=$(openai list-messages --thread "$THREAD_ID" --order "asc" --after "$MESSAGE_ID")
    has_more=$(jq -r '.has_more' <<< "$messages")

    IFS=$'\n' read -r -d '' -a lines < <(jq -c '.data[] | {id: .id, role: .role, text: .content[0].text.value}' <<< "$messages" && printf '\0')

    for line in "${lines[@]}"; do
      MESSAGE_ID=$(jq -r '.id' <<< "$line")
      role=$(jq -r '.role' <<< "$line")
      text=$(jq -r '.text' <<< "$line")

      # Call the callback function with necessary parameters
      "$callback" "$role" "$text"
    done
  done
}

# ------------------------------------------------------------------------------
# Feature: ASSISTANTS
# ------------------------------------------------------------------------------
list-models() {
  curl -s "https://api.openai.com/v1/models" \
    -H "Authorization: Bearer $OPENAI_API_KEY" \
    | jq -r '.data.[].id' \
    | grep gpt \
    | sort
}

select-model() {
  local model="${1:-"$OPENAI_MODEL"}"
  local models=()

  while IFS=':' read -r name; do
    models+=("$name")
  done < <(list-models)

  gum choose --header "Select a model" --selected "$model" "${models[@]}" || return 0
}

# ------------------------------------------------------------------------------
# Creates a new custom GPT model and emits its id.
# ------------------------------------------------------------------------------
add-assistant() {
  local name
  local instructions
  local description
  local model

  name=$(gum input --prompt 'Enter a name for the assistant: ' --placeholder 'G. Peetee') || return 0

  [[ -z "$name" ]] && {
    name="G. Peetee"
  }

  instructions=$(
    gum write \
      --width 100 \
      --height 30 \
      --char-limit 0 \
      --placeholder "You are a helpful assistant" \
      --header "GPT instructions" \
      --header.foreground="#66FF66" \
      --cursor.mode=static \
      || true

  )

  [[ -z "$instructions" ]] && {
    instructions="You are a helpful chatbot."
  }

  description=$(gum input --prompt 'Enter an optional description: ' --placeholder 'Gene Peetee, the helpful chatbot.')

  # shellcheck disable=SC2119
  model=$(select-model)

  ASSISTANT_ID=$(create-assistant "$name" "$instructions" "$description" "$model")
  ASSISTANT_NAME="$name"

  save-selected-assistant
}

# ------------------------------------------------------------------------------
# Lists the assistants and prompts the user to select one. Sets the global
# ASSISTANT_ID and ASSISTANT_NAME variables to the selected assistant.
# ------------------------------------------------------------------------------
edit-assistant() {
  local result

  local name
  local instructions
  local description
  local model

  local new_name
  local new_instructions
  local new_description
  local new_model

  if ! result=$(openai get-assistant --assistant "$ASSISTANT_ID"); then
    warn "error getting assistant"
    die "$result"
  fi

  name=$(jq -r '.name' <<< "$result")
  instructions=$(jq -r '.instructions' <<< "$result")
  description=$(jq -r '.description' <<< "$result")
  model=$(jq -r '.model' <<< "$result")

  if ! new_name=$(gum input --prompt 'Enter a new name: ' --value "$name"); then
    return
  fi

  if ! new_instructions=$(gum write --width 100 --height 30 --char-limit 0 --value "$instructions" --header "Modify GPT instructions" --header.foreground="#66FF66" --cursor.mode=static); then
    return
  fi

  if ! new_description=$(gum input --prompt 'Modify description: ' --value "$description"); then
    return
  fi

  if ! new_model=$(select-model "$model"); then
    return
  fi

  if [ -z "$new_name" ] || [ -z "$new_instructions" ]; then
    # The user pressed escape on name or instructions, so we skip updating the
    # description.
    new_name=$name
    new_instructions=$instructions
  fi

  ASSISTANT_ID=$(modify-assistant "$new_name" "$new_instructions" "$new_description" "$new_model")
  ASSISTANT_NAME="$new_name"

  save-selected-assistant
}

# ------------------------------------------------------------------------------
# Deletes the selected assistant. I believe this cascades to delete all
# associated threads and messages but I'm not sure.
# ------------------------------------------------------------------------------
del-assistant() {
  local old_id="$ASSISTANT_ID"
  local tmp

  gum confirm "Delete assistant $ASSISTANT_NAME (id: $ASSISTANT_ID)?" || return 0
  delete-assistant || return 0

  tmp=$(mktemp)
  if grep -v "^$old_id:" "$CHAT_HISTORY_FILE" > "$tmp"; then
    mv "$CHAT_HISTORY_FILE" "$CHAT_HISTORY_FILE.bak"
    mv "$tmp" "$CHAT_HISTORY_FILE"
    rm "$CHAT_HISTORY_FILE.bak"
  fi
}

# ------------------------------------------------------------------------------
# Saves the currently selected assistant to a config file so that it can be
# restored on the next run.
# ------------------------------------------------------------------------------
save-selected-assistant() {
  echo "$ASSISTANT_ID" > "$LAST_ASSISTANT_FILE"
  echo "$ASSISTANT_NAME" >> "$LAST_ASSISTANT_FILE"
}

# ------------------------------------------------------------------------------
# Restores the last assistant saved by `save-selected-assistant`. If the file
# does not exist or is malformed, it is deleted and a new one is created.
# ------------------------------------------------------------------------------
restore-selected-assistant() {
  local last_assistant

  if [ -f "$LAST_ASSISTANT_FILE" ]; then
    readarray -t last_assistant < "$LAST_ASSISTANT_FILE"

    if [ "${#last_assistant[@]}" -ne 2 ]; then
      rm "$LAST_ASSISTANT_FILE"
      touch "$LAST_ASSISTANT_FILE"
      return 0
    fi

    ASSISTANT_ID="${last_assistant[0]}"
    ASSISTANT_NAME="${last_assistant[1]}"
  fi
}

# ------------------------------------------------------------------------------
# Feature: CHAT
# ------------------------------------------------------------------------------

# ------------------------------------------------------------------------------
# Files and command outputs can be too long to send in a single message, so
# they are broken up into multiple messages. This function splits the string
# into chunks based on characters, then prints each chunk sperated by a null
# byte.
# ------------------------------------------------------------------------------
split-into-chunks() {
  local chunk_size=$1
  local str=$2
  local chunk

  while [ -n "$str" ]; do
    chunk=${str:0:$chunk_size}
    echo -n "$chunk" # Print chunk without a trailing newline
    echo -ne "\0"    # Use null byte as delimiter
    str=${str:$chunk_size}
  done
}

# ------------------------------------------------------------------------------
# Asks the user to select a file with autocomplete.
# ------------------------------------------------------------------------------
select-file() {
  find . -type f | grep -v '.git' | gum filter --header "$1" --fuzzy
}

# ------------------------------------------------------------------------------
# Reads the file contents and prints them out. If there is an error, such as
# a permissions issue or an invalid file path, the user is prompted to either
# skip the file, enter a new file path, or try again.
#
# Prints out the file path, followed by a string delimited "###FILEPATH###",
# then the file's contents.
# ------------------------------------------------------------------------------
readfile-for-attachment() {
  local filepath="$1"
  local content
  local retry=true
  local choice
  local error_message

  while [ "$retry" = true ]; do
    retry=false

    if content=$(cat "$filepath" 2>&1); then
      printf "%s###FILEPATH###%s" "$filepath" "$content"
      return 0
    else
      error_message="$content" # Capture the error message
      choice=$(gum choose --header "Error reading file: $error_message" "Skip this file" "Enter a new file path" "Try again")

      case "$choice" in
        "Skip this file")
          return 1 # Indicate to skip this file
          ;;
        "Enter a new file path")
          filepath=$(select-file "Enter a new file path")
          retry=true
          ;;
        "Try again")
          retry=true
          ;;
      esac
    fi
  done
}

# ------------------------------------------------------------------------------
# Attaches a file to the current message. If no file is specified, the user is
# prompted to select one. If the file is too large, the user is prompted to
# either truncate the file or split it into multiple messages. Regardless of
# how many "messages" this function generates, it prints each one
# base64-encoded to make it easier for the caller to read the output, because
# bash makes it sooooooooo easy to return multiple strings from a function.
# ------------------------------------------------------------------------------
attach-file() {
  local filepath
  local content
  local action
  local msg

  if [ -z "$*" ]; then
    filepath=$(select-file "Select a file to attach") || return 0
  else
    filepath="$1"
  fi

  # Remove leading whitespace, if any
  filepath="${filepath#"${filepath%%[![:space:]]*}"}"

  # Remove trailing whitespace, if any
  filepath="${filepath%"${filepath##*[![:space:]]}"}"

  # Read the file contents, giving the user an opportunity to handle file read
  # errors.
  file_info=$(readfile-for-attachment "$filepath") || return $?

  separator="###FILEPATH###"
  filepath="${file_info%%"$separator"*}"
  content="${file_info#*"$separator"}"
  content_length=$(wc -m <<< "$content")

  if ((content_length > MAX_FILE_SIZE)); then
    warn "The file content ($content_length) exceeds the maximum allowed size ($MAX_FILE_SIZE)."

    action_truncate="Truncate file to fit the limit"
    action_split="Split the content into multiple messages"
    action_cancel="Cancel the operation"

    action=$(gum choose --header "How would you like to proceed?" "$action_truncate" "$action_split" "$action_cancel")

    case "$action" in
      "$action_truncate")
        content="${content:0:MAX_FILE_SIZE}"
        ;;

      "$action_split")
        # Calculate maximum size for a chunk of content, considering markdown and headers
        local total_characters=${#content}
        local part_no=0
        local chunk
        local pct

        while IFS= read -r -d '' chunk; do # Read until null byte
          part_no=$((part_no + 1))

          pct=$((part_no * MAX_FILE_SIZE * 100 / total_characters))
          if ((pct > 100)); then
            pct=100
          fi

          info "Sending chunk ${part_no} (${#chunk} characters) - ${pct}% complete"

          # shellcheck disable=SC2016
          msg=$(printf 'Attached file content: `%s` (part %s)\n```\n%s\n```\n' "$filepath" "$part_no" "$chunk")
          echo -e "$msg" | base64
        done < <(split-into-chunks "$MAX_FILE_SIZE" "$content")

        return 0
        ;;

      "$action_cancel")
        return 0
        ;;

      *)
        return 0
        ;;
    esac
  fi

  # shellcheck disable=SC2016
  msg=$(printf 'Attached file content: `%s`\n```\n%s\n```\n' "$filepath" "$content")
  echo -e "$msg" | base64
}

# ------------------------------------------------------------------------------
# Attaches the output of a command to the current message. If the command's
# output is too large, it is split up into multiple messages. Just like
# `attach-file`, it prints each message base64-encoded to make it easier for
# the caller to read the output.
# ------------------------------------------------------------------------------
attach-command-output() {
  local command="$1"
  local command_output
  local command_status
  local msg

  # Execute the command and capture its output
  command_output=$(eval "$command" 2>&1)
  command_status=$?
  command_output=$(printf "\`\`\`\n\$ %s\n%s\n# exit status: %s\n\`\`\`" "$command" "$command_output" "$command_status")

  if ((${#command_output} > MAX_MSG_SIZE)); then
    local total_characters=${#command_output}
    local part_no=0
    local chunk
    local pct

    echo -e "Output of \`$command\` is too large to send in a single message. Sending in multiple messages instead." | base64

    while IFS= read -r -d '' chunk; do # Read until null byte
      part_no=$((part_no + 1))

      pct=$((part_no * MAX_FILE_SIZE * 100 / total_characters))
      if ((pct > 100)); then
        pct=100
      fi

      # shellcheck disable=SC2016
      msg=$(printf 'Command `%s` output continued (part %s):\n```\n%s\n```' "$command" "$part_no" "$chunk")
      echo -e "$msg" | base64
    done < <(split-into-chunks "$MAX_MSG_SIZE" "$command_output")

    return 0
  fi

  # shellcheck disable=SC2016
  msg=$(printf 'Command `%s` output:\n```\n%s\n```' "$command" "$command_output")
  echo -e "$msg" | base64
}

# ------------------------------------------------------------------------------
# Attaches the content of a URL to the current message. If the content is too
# large, it is split up into multiple messages.
# ------------------------------------------------------------------------------
# shellcheck disable=SC2317
attach-web-content() {
  local url="$1"
  local content
  local msg

  # Remove leading and trailing whitespace
  url="${url#"${url%%[![:space:]]*}"}"
  url="${url%"${url##*[![:space:]]}"}"

  info "Fetching $url..."

  if ! content=$(curl -s "$url"); then
    warn "error fetching url: $url"
    return 0
  fi

  if ((${#content} > MAX_MSG_SIZE)); then
    local total_characters=${#content}
    local part_no=0
    local chunk
    local pct

    echo -e "The contents of \`$url\` are too large to send in a single message. Sending in multiple messages instead." | base64

    while IFS= read -r -d '' chunk; do # Read until null byte
      part_no=$((part_no + 1))

      pct=$((part_no * MAX_FILE_SIZE * 100 / total_characters))
      if ((pct > 100)); then
        pct=100
      fi

      # shellcheck disable=SC2016
      msg=$(printf 'Attached web content: `%s` (continued - part %s):\n```\n%s\n```' "$url" "$part_no" "$chunk")
      echo -e "$msg" | base64

      sleep 1 # Adding a short delay to avoid rapid message sending
    done < <(split-into-chunks "$MAX_MSG_SIZE" "$content")

    return 0
  fi

  # shellcheck disable=SC2016
  msg=$(printf 'Attached web content: `%s`\n```\n%s\n```\n' "$url" "$content")
  echo -e "$msg" | base64
}

# ------------------------------------------------------------------------------
# Chat supports a number of embedded directives, like \f to attach a file, and
# \exec to include the output of a command (or \begin + \end for a multi-line
# command). This function parses the message into text, exec, and file
# sections, and prints each section out preceded by a label (TEXT, FILE, or
# EXEC).
# ------------------------------------------------------------------------------
parse-message() {
  local message="$1"
  local is_in_code_fence=false
  local command_block=""
  local inside_command_block=false
  local parsed_lines=()

  while IFS= read -r line; do
    # Do not process any directives within a markdown code fence
    if [[ "$line" =~ ^\`\`\` ]]; then
      # Toggle the is_in_code_fence flag
      if [ "$is_in_code_fence" = true ]; then
        is_in_code_fence=false
      else
        is_in_code_fence=true
      fi

      parsed_lines+=("TEXT $line")

    # We're not in a code fence - GAME ON
    elif [ "$is_in_code_fence" != true ]; then
      # ------------------------------------------------------------------------
      # Command block
      # ------------------------------------------------------------------------
      # Beginning of command block
      if [[ "$line" == \\begin ]]; then
        inside_command_block=true
        command_block=""

      # End of command block
      elif [[ "$line" == \\end ]] && $inside_command_block; then
        inside_command_block=false

        if [[ $command_block =~ [^[:space:]] ]]; then
          parsed_lines+=("EXEC $command_block")
        fi

        command_block=""

      # Inside of command block - collect commands
      elif $inside_command_block; then
        if [[ "$line" =~ \\$ ]]; then
          # Append part of command, trimming the trailing backslash
          command_block+="${line%\\} "
        else
          command_block+="$line"

          if [[ $command_block =~ [^[:space:]] ]]; then
            parsed_lines+=("EXEC $command_block")
          fi

          command_block=""
        fi

      # ------------------------------------------------------------------------
      # Single line command
      # ------------------------------------------------------------------------
      elif [[ "$line" =~ ^\\exec\ (.*) ]]; then
        parsed_lines+=("EXEC ${BASH_REMATCH[1]}")

      # ------------------------------------------------------------------------
      # File block
      # ------------------------------------------------------------------------
      elif [[ "$line" =~ ^\\f(.*) ]]; then
        parsed_lines+=("FILE ${BASH_REMATCH[1]}")

      # ------------------------------------------------------------------------
      # Web block
      # ------------------------------------------------------------------------
      elif [[ "$line" =~ ^\\www(.*) ]]; then
        parsed_lines+=("WWW ${BASH_REMATCH[1]}")

      # ------------------------------------------------------------------------
      # Normal message
      # ------------------------------------------------------------------------
      else
        parsed_lines+=("TEXT $line")
      fi
    else
      parsed_lines+=("TEXT $line")
    fi
  done <<< "$message"

  # Add any remaining non-empty commands to parsed_lines
  if [[ -n $command_block && $command_block =~ [^[:space:]] ]]; then
    parsed_lines+=("EXEC $command_block")
  fi

  printf "%s\n" "${parsed_lines[@]}"
}

# ------------------------------------------------------------------------------
# A little helper function that joins an array of strings with newlines.
# ------------------------------------------------------------------------------
join-by-newline() {
  local IFS=$'\n'
  echo "$*"
}

# ------------------------------------------------------------------------------
# This function accepts the output of `parse-message` and sends each one as a
# separate message to the current thread. Since a message can have multiple
# directives embedded within them, `parse-message` processes the entire message
# line by line, attaching a label to each line. This function iterates over
# those labeled lines (grouping subsequent TEXT lines together) and sends each
# as a separate message to the current thread, with a short sleep between each
# as a naive rate limiter.
# ------------------------------------------------------------------------------
handle-parsed-message() {
  local parsed_lines=("$@")
  local messages=()
  local text_content=()

  for line in "${parsed_lines[@]}"; do
    case "$line" in
      TEXT*)
        text_content+=("${line#TEXT }")
        ;;

      WWW*)
        if [ "${#text_content[@]}" -gt 0 ]; then
          messages+=("$(join-by-newline "${text_content[@]}")")
          text_content=()
        fi

        local url="${line#WWW }"
        local msgs=()
        local msg
        readarray -t msgs < <(attach-web-content "$url")

        for msg in "${msgs[@]}"; do
          msg=$(base64 --decode <<< "$msg")
          if [ -n "$msg" ]; then
            messages+=("$msg")
          fi
        done
        ;;

      EXEC*)
        if [ "${#text_content[@]}" -gt 0 ]; then
          messages+=("$(join-by-newline "${text_content[@]}")")
          text_content=()
        fi

        local exec_line="${line#EXEC }"
        local msgs=()
        local msg

        info "Running command:"
        info "    $ ${line#EXEC }"
        readarray -t msgs < <(attach-command-output "$exec_line")
        info "✓"

        for msg in "${msgs[@]}"; do
          msg=$(base64 --decode <<< "$msg")
          if [ -n "$msg" ]; then
            messages+=("$msg")
          fi
        done
        ;;

      FILE*)
        if [ "${#text_content[@]}" -gt 0 ]; then
          messages+=("$(join-by-newline "${text_content[@]}")")
          text_content=()
        fi

        local file_path="${line#FILE }"
        local msgs=()
        local msg
        if readarray -t msgs < <(attach-file "${file_path}"); then
          for msg in "${msgs[@]}"; do
            msg=$(base64 --decode <<< "$msg")
            if [ -n "$msg" ]; then
              messages+=("$msg")
            fi
          done
        else
          warn "Failed to attach file: ${file_path}. Skipping file attachment."
        fi
        ;;

      *)
        warn "Unhandled message type: $line"
        ;;
    esac
  done

  if [ "${#text_content[@]}" -gt 0 ]; then
    messages+=("$(join-by-newline "${text_content[@]}")")
    text_content=()
  fi

  for msg in "${messages[@]}"; do
    send-msg "$msg"
    sleep 0.5
  done

  echo "${#messages[@]}"
}

# ------------------------------------------------------------------------------
# Prints out a formatted message from the user.
# ------------------------------------------------------------------------------
# Called indirectly in print-message-callback
# shellcheck disable=SC2317
print-user-msg() {
  local msg="$1"

  if [[ "$msg" =~ ^Attached\ file\ content: ]]; then
    msg=$(echo "$msg" | sed -n '1p')
  fi

  if [[ "$msg" =~ ^Attached\ web\ content: ]]; then
    msg=$(echo "$msg" | sed -n '1p')
  fi

  echo
  gum style --foreground '#66FF66' --bold 'You:'
  gum format -- "$msg"
}

# ------------------------------------------------------------------------------
# Prints out a formatted message from GPT with correct handling of initial dashes.
# ------------------------------------------------------------------------------
# Called indirectly in print-message-callback
# shellcheck disable=SC2317
print-gpt-msg() {
  local msg="${1:-}"

  gum style --foreground '#66CCFF' --bold "$ASSISTANT_NAME:"

  if [ -n "$msg" ]; then
    echo
    gum format -- "$msg"
  fi
}

# ------------------------------------------------------------------------------
# This is the callback `print-messages` passes to `list-messages`. This
# function is responsible for deciding how to format each message based on the
# role, using either `print-user-msg` or `print-gpt-msg`.
# ------------------------------------------------------------------------------
# shellcheck disable=SC2317
print-message-callback() {
  local role=$1
  local text=$2

  if [[ "$role" == "user" ]]; then
    print-user-msg "$text"
  else
    print-gpt-msg "$text"
  fi
}

# ------------------------------------------------------------------------------
# Retrieves the messages from the current thread and prints them out formatted.
# Accepts two optional parametes, `limit`, and `get_all`, which are passed to
# `list-messages`. See `list-messages` for detail.
# ------------------------------------------------------------------------------
print-messages() {
  list-messages print-message-callback >> "$CONVERSATION_TMP"
  clear
  cat "$CONVERSATION_TMP"
}

# ------------------------------------------------------------------------------
# This is the main chat loop. It retrieves the messages from the current thread
# and prints them out formatted, then prompts the user for a message to send.
#
# Note that there are a few unadvertised directives that are not handled by
# `parse-message` that are handled directly here. They are mainly for testing.
# ------------------------------------------------------------------------------
# shellcheck disable=SC2016
CHAT_COMMANDS=(
  '\f $filename'
  '\www $url'
  '\exec $cmd'
  '\begin $cmd \end'
  '\run'
)

chat() {
  local placeholder
  local needs_summary=false
  local initialized=false
  local msgs_sent
  local msg

  placeholder="ctrl-d submits | esc cancels$(printf " | %s" "${CHAT_COMMANDS[@]}")"

  # Come on, shellcheck, this is clearly a string, not an array. Seriously,
  # `read` operates just fine on here-strings.
  # shellcheck disable=SC2178
  local messages=""

  if [ -z "$ASSISTANT_ID" ]; then
    die "No assistant is selected"
  fi

  if [ -z "$THREAD_ID" ]; then
    die "No chat is selected"
  fi

  if [ "$THREAD_LABEL" == "New chat with $ASSISTANT_NAME" ]; then
    needs_summary=true
  fi

  while true; do
    if [ "$initialized" == false ]; then
      if [ "$THREAD_LABEL" != "New chat with $ASSISTANT_NAME" ]; then
        info "Retrieving conversation..."
        move-chat-to-beginning-of-history

        clear
      fi

      initialized=true
    fi

    print-messages

    while true; do
      msg=$(
        gum write \
          --width 140 \
          --height 15 \
          --char-limit 0 \
          --placeholder "$placeholder" \
          --header ">>> $THREAD_LABEL (id: $THREAD_ID) <<<" \
          --header.foreground="#66FF66" \
          --cursor.mode=static \
          || true
      )

      case "$msg" in
        '')
          break 2
          ;;

        '\rename')
          save-to-chat-history
          break
          ;;

        '\run')
          info "Requesting a new response to the conversation..."

          get-response

          if [ "$needs_summary" = true ]; then
            save-to-chat-history
            needs_summary=false
          fi

          break
          ;;

        *)
          IFS=$'\n' read -r -d '' -a parsed_lines < <(parse-message "$msg" && printf '\0')
          msgs_sent=$(handle-parsed-message "${parsed_lines[@]}")

          if [ "$msgs_sent" == 0 ]; then
            break
          fi

          # If any messages were sent, update the display to show them while we
          # wait for the thread to run.
          print-messages

          # Stream GPT's output to the terminal. On the next loop, it will
          # clear the screen and print the updated, formatted conversation.
          get-response

          if [ "$needs_summary" = true ]; then
            save-to-chat-history
            needs_summary=false
          fi

          break
          ;;

      esac
    done
  done
}

# ------------------------------------------------------------------------------
# Whenever a existing thread is opened, this function is used to push it to the
# top of the chat history file so that chat history is displayed in most recent
# order.
# ------------------------------------------------------------------------------
move-chat-to-beginning-of-history() {
  local temp_file
  local chat_entry

  temp_file=$(mktemp)

  # Search for the chat entry
  chat_entry=$(grep "^$ASSISTANT_ID:$THREAD_ID:" "$CHAT_HISTORY_FILE" || true)

  # Remove the existing chat entry, if any, from the chat history.
  # Use '|| true' to prevent grep's exit status from causing the script to exit.
  grep -v "^$ASSISTANT_ID:$THREAD_ID:" "$CHAT_HISTORY_FILE" > "$temp_file" || true

  # Add the chat entry to the "beginning" of the history file if it exists
  if [ -n "$chat_entry" ]; then
    {
      echo "$chat_entry"
      cat "$temp_file"
    } > "$CHAT_HISTORY_FILE"
  else
    # If no entry is found, just copy temp_file back
    cat "$temp_file" > "$CHAT_HISTORY_FILE"
  fi

  # Clean up
  rm -f "$temp_file"
}

# ------------------------------------------------------------------------------
# When a chat is first started, this function is used to create save it to the
# chat history file, adding a GPT-3.5-turbo powered summary to it.
# ------------------------------------------------------------------------------
save-to-chat-history() {
  # goddamnit apple
  if [[ "$(uname)" == "Darwin" ]]; then
    SED_INPLACE_FLAG='-i ""'
  else
    SED_INPLACE_FLAG='-i'
  fi

  if grep "$ASSISTANT_ID:$THREAD_ID:" "$CHAT_HISTORY_FILE" > /dev/null; then
    eval "sed $SED_INPLACE_FLAG '/$ASSISTANT_ID:$THREAD_ID:/d' '$CHAT_HISTORY_FILE'"
  fi

  THREAD_LABEL=$(summarize-chat)
  printf '%s\n' "$ASSISTANT_ID:$THREAD_ID:$THREAD_LABEL" >> "$CHAT_HISTORY_FILE"
}

# ------------------------------------------------------------------------------
# Uses the "traditional" completion API to generate a summary of a chat based
# on the first few messages.
# ------------------------------------------------------------------------------
summarize-chat() {
  get-completion \
    'gpt-3.5-turbo-16k' \
    'You are a bot that summarizes conversations tersely' \
    "$(openai list-messages --thread "$THREAD_ID" --order asc)\n\nSummarize the above conversation in 10 words or less"
}

# ------------------------------------------------------------------------------
# Creates a new thread, sets the global THREAD_ID and THREAD_LABEL variables to
# that of the newly created thread, and then passes the buck to `chat`.
# ------------------------------------------------------------------------------
new-chat() {
  THREAD_ID=$(start-thread)
  THREAD_LABEL="New chat with $ASSISTANT_NAME"
  MESSAGE_ID=""
  echo "" > "$CONVERSATION_TMP"
  chat
}

# ------------------------------------------------------------------------------
# Prints out the list of previously saved chats for the current assistant. Used
# in the `chat-history-menu`.
# ------------------------------------------------------------------------------
list-chats() {
  local assistant_id
  local thread_id
  local label

  if [ -z "$ASSISTANT_ID" ]; then
    die "No assistant is selected"
  fi

  tac "$CHAT_HISTORY_FILE" | while IFS=':' read -r assistant_id thread_id label; do
    if [ "$assistant_id" == "$ASSISTANT_ID" ]; then
      echo -e "$label ($(gum style --faint "id: $thread_id"))"
    fi
  done
}

# ------------------------------------------------------------------------------
# Menus
# ------------------------------------------------------------------------------
assistant-menu() {
  local choices=()
  local header
  local selected

  while true; do
    clear
    choices=()

    while IFS=':' read -r id name; do
      choices+=("$name ($(gum style --faint "id: $id"))")
    done < <(list-assistants)

    choices+=("Create a new assistant")

    if [ -n "$ASSISTANT_ID" ]; then
      choices+=('Edit' 'Delete')
    fi

    choices+=("Back")

    header="Assistants"
    if [ -n "$ASSISTANT_ID" ]; then
      header="$header (selected: $ASSISTANT_NAME)"
    fi

    selected=$(gum choose --header "$header" "${choices[@]}") || return 0

    case "$selected" in
      "" | "Back")
        break
        ;;

      "Create a new assistant")
        add-assistant
        ;;

      "Edit")
        edit-assistant
        ;;

      "Delete")
        del-assistant
        ;;

      *)
        ASSISTANT_ID="$(echo "$selected" | sed -E 's/.*id: ([^)]+).*/\1/')"
        ASSISTANT_NAME="$(echo "$selected" | sed -E 's/^(.*) \(id: [^)]+\)$/\1/')"
        save-selected-assistant
        break
        ;;
    esac
  done
}

chat-menu() {
  local choices=()
  local selected

  if [ -z "$ASSISTANT_ID" ]; then
    die "No assistant is selected"
  fi

  while true; do
    clear

    choices=("Start a new chat")

    if grep "^$ASSISTANT_ID:" "$CHAT_HISTORY_FILE" > /dev/null; then
      choices+=("Continue a previous chat with $ASSISTANT_NAME")
    fi

    choices+=("Back")

    selected=$(gum choose --header "Chats" "${choices[@]}") || true

    case "$selected" in
      "" | "Back")
        break
        ;;

      "Start a new chat")
        new-chat
        ;;

      "Continue a previous chat with $ASSISTANT_NAME")
        chat-history-menu
        ;;
    esac
  done
}

chat-history-menu() {
  local choice
  local height
  local chats=()

  while true; do
    clear
    height=$(($(tput lines) - 2)) # -2 lines for gum filter's header and input

    chats=()
    while IFS= read -r line; do
      chats+=("$line")
    done < <(list-chats)

    if [ "${#chats[@]}" -eq 0 ]; then
      info 'No chat history found'
      read -r -p "Press enter to continue..." || true
      return 0
    fi

    choice=$(
      printf "%s\n" "${chats[@]}" \
        | gum filter \
          --header "Conversation history" \
          --placeholder "Search" \
          --no-fuzzy \
          --height "$height" \
        || true
    )

    if [ -n "$choice" ]; then
      THREAD_ID="$(echo "$choice" | sed -E 's/.*id: ([^)]+).*/\1/')"
      THREAD_LABEL="$(echo "$choice" | sed -E 's/^(.*) \(id: [^)]+\)$/\1/')"
      MESSAGE_ID=""
      echo "" > "$CONVERSATION_TMP"
      chat
    else
      break
    fi
  done
}

main-menu() {
  local choices
  local selected

  while true; do
    clear

    choices=()
    if [ -n "$ASSISTANT_NAME" ]; then
      choices+=("Chat with $ASSISTANT_NAME")
    fi

    choices+=(
      "Assistants"
      "Quit"
    )

    selected=$(gum choose --header "What would you like to do?" "${choices[@]}") \
      || return 0

    case "$selected" in
      "")
        continue
        ;;

      "Assistants")
        assistant-menu
        ;;

      "Chat with $ASSISTANT_NAME")
        chat-menu
        ;;

      "Quit")
        return 0
        ;;
    esac
  done
}

main() {
  validate-env
  main-menu
  exit 0
}

main