# aish: AI shell completions
# <https://github.com/chr15m/aish>

# vi: ft=bash
# Why no hash-bang?
# It's so this script can be sourced under both bash and zsh.

# Website: <https://mccormick.cx>
# Donate:  <https://www.patreon.com/chr15m>
# Donate:  <ko-fi.com/chr15m>

# TODO:
# - set up basic tests via github actions

if [ "$1" = "" -o "$1" = "-h" ]
then
  echo "Usage: . aish QUERY"
  echo
  echo "QUERY:   tell the AI what you want"
  echo "Example: '. aish a for loop going from 5 to 15'."
  echo "Note:    You must include the . to source the script."
  echo "         The result will be available from history with the up key."
  echo
  echo 'You an add variables to ~/.aish to configure it:'
  echo 'OPENAI_API_KEY="sk-..."'
  echo 'AISH_MODEL=gpt-4 # defaults to gpt-3.5-turbo (gpt-4 is recommended)'
  echo 'AISH_URL=http... # defaults to https://api.openai.com - set to any OpenAI compatible API.'
  echo ''
  echo '- Create a new API key here: <https://platform.openai.com/api-keys>'
  echo '- List the available models with `aish models`'
  echo '- Requests will be logged to ~/.aish-log'
  exit 0
fi

# source aish config file with overrides if it exists
if [ -f ~/.aish ]
then
  . ~/.aish
fi

AISH_URL=${AISH_URL:-https://api.openai.com}
BEARER_TOKEN=${AISH_BEARER_TOKEN:-}
API_KEY=${AISH_API_KEY:-${OPENAI_API_KEY:-}}

# check if API credentials are set
if [ -z "$BEARER_TOKEN" ] && [ -z "$API_KEY" ]
then
  echo "Error: no API credentials configured."
  echo
  echo "Set AISH_BEARER_TOKEN for bearer authentication, or set"
  echo "AISH_API_KEY / OPENAI_API_KEY for basic authentication."
  echo
  echo "For OpenAI, create a key at"
  echo "<https://platform.openai.com/account/api-keys>."
  return 1
fi

if [ "$1" = "models" ]
then
  echo "Fetching available model list from the API."
  if [ -n "$BEARER_TOKEN" ]
  then
    curl -s -H "Authorization: Bearer $BEARER_TOKEN" \
      ${AISH_URL}/v1/models |
      grep gpt- | cut -d'"' -f4 | sort -r
  else
    curl -s -u ":$API_KEY" ${AISH_URL}/v1/models |
      grep gpt- | cut -d'"' -f4 | sort -r
  fi
  exit 0
fi

# test if user has sourced the script and notify them if not

_throw_source_error () {
  echo "Please source aish in a bash or zsh shell like this:"
  echo ". aish ${@}"
  # echo ${1}
  exit 1
}

OS=`uname`
OS="${OS//Darwin/Mac OSX}"

case $BASH_VERSION in *.*)
  SHELLNAME="bash"
  ;;
esac
case $ZSH_VERSION  in *.*)
  SHELLNAME="zsh"
  ;;
esac
case "$VERSION" in *zsh*)
  SHELLNAME="zsh"
  ;;
esac

MODEL=${AISH_MODEL:-gpt-3.5-turbo}

#echo "SHELLNAME: ${SHELLNAME}"
#echo "OS: ${OS}"
#echo "MODEL: ${MODEL}"

if [ "${SHELLNAME}" != "zsh" ] && [ "$SHELLNAME" != "bash" ]
then
  _throw_source_error "${@}"
fi

if [ "${SHELLNAME}" = "bash" ]
then
  (return 0 2>/dev/null) && sourced=1 || sourced=0
  if [ $sourced -eq 0 ]; then
    _throw_source_error "${@}"
  fi
fi

if [ "${SHELLNAME}" = "zsh" ]
then
  if [[ "${ZSH_EVAL_CONTEXT}" = "toplevel" ]]
  then
    _throw_source_error "${@}"
  fi
  if ! (( zsh_eval_context[(I)file] )); then
    _throw_source_error "${@}"
  fi
fi

# exit from ctrl-C without exiting user's shell
trap "return 0" SIGINT

AISH_SHELL="${SHELLNAME}"
AISH_OS="${OS}"

# start the actual request to OpenAI

request=`echo "${@}" | sed -r 's/"/\\\"/g'`

data='{
  "model": "'${MODEL}'",
  "messages": [
  {"role": "system", "content": "You are an AI shell-scripting expert called Aish."},{"role": "user", "content": "Please write a one-liner for '${AISH_SHELL}' shell for performing the following task. Assume you have access to all of the commonly available unix commands on a modern '${AISH_OS}' system. You can separate multi-line commands with a semicolon or double ampersand but please put them on one line. Respond without using markdown formatting. Please only return the exact bash one-liner command without any superfluous words explaining it, and no formatting. Do NOT use backticks. Do not put backticks around the one-liner. Code should be bare without any markdown formatting. IMPORTANT: If you need to ask for clarification or provide any information other than a one-liner solution, please prefix your response with the exact phrase AISH_FEEDBACK: in all caps, with a colon. If you are responding with anything other than a shell script one-liner you MUST include AISH_FEEDBACK: (including the colon) as the start of the response.\n\n Here is the task I would like you to carry out with the one-liner: '$request'"}
  ]
}'

# echo "Sending: ${data}"

if [ -n "$BEARER_TOKEN" ]
then
  json=$( curl -s ${AISH_URL}/v1/chat/completions \
    -H "Authorization: Bearer ${BEARER_TOKEN}" \
    -H 'Content-Type: application/json' \
    -d "${data}" )
else
  json=$( curl -s ${AISH_URL}/v1/chat/completions \
    -u ":${API_KEY}" \
    -H 'Content-Type: application/json' \
    -d "${data}" )
fi

if printf '%s' "${json}" | grep -q '"content":'
then
  now=`date "+%Y-%m-%dT%H:%M:%S%:z"`
  echo "$now ${request}" >> ~/.aish-log
  if command -v python3 >/dev/null 2>&1
  then
    result=$(printf '%s' "${json}" | python3 -c 'import json
import sys

data = json.load(sys.stdin)
choices = data.get("choices") or []
if not choices:
    sys.exit(1)
message = choices[0].get("message") or {}
content = message.get("content")
if content is None:
    sys.exit(1)
sys.stdout.write(content)')
    status=$?
  else
    result=""
    status=1
  fi
  if [ $status -ne 0 ] || [ -z "$result" ]
  then
    result=$(echo "${json}" |
      sed -nr 's/"content":[[:space:]]*"(.*)"\,/\1/p' |
      sed -e 's/^[ \t]*//' -e 's/\\"/"/g' | head -n 1)
  fi
  if [ -z "$result" ]
  then
    echo "${json}"
    return 1
  fi
  if [ "${result#AISH_FEEDBACK: }" != "$result" ]
  then
    echo "${result#AISH_FEEDBACK: }"
  else
    echo "Hit Enter to add this command to history or ctrl-C to abort:"
    echo -n "$ ${result}"
    read -r _
    if [ "${SHELLNAME}" = "zsh" ]
    then
      print -s "${result}"
    fi
    if [ "${SHELLNAME}" = "bash" ]
    then
      set -o history
      shopt -s histappend
      history -s "${result}"
      history -w
    fi
    echo "Press the up arrow to access the solution."
  fi
else
  echo "${json}"
fi