#!/usr/bin/env bash

set -eu

if [ $# -ge 1 ] && [ "${1:0:2}" == "-h" ]; then
    echo "Usage: update.sh <release>"
    echo "Updates InvokeAI to use the indicated version of the code base."
    echo "Find the version or branch for the release you want, and pass it as the argument."
    echo "For example: update.sh v2.2.5 for release 2.2.5."
    echo "             update.sh main for the current development version."
    echo ""
    echo "If no argument provided then will install the version tagged with 'latest', equivalent to"
    echo "update.sh latest"
    exit -1
fi

INVOKE_AI_VERSION=${1:-latest}

INVOKE_AI_SRC="https://github.com/invoke-ai/InvokeAI/archive/$INVOKE_AI_VERSION.zip"
INVOKE_AI_DEP=https://raw.githubusercontent.com/invoke-ai/InvokeAI/$INVOKE_AI_VERSION/environments-and-requirements/requirements-base.txt
INVOKE_AI_MODELS=https://raw.githubusercontent.com/invoke-ai/InvokeAI/$INVOKE_AI_VERSION/configs/INITIAL_MODELS.yaml

# ensure we're in the correct folder in case user's CWD is somewhere else
scriptdir=$(dirname "$0")
cd "$scriptdir"

function _err_exit {
    if test "$1" -ne 0
    then
	echo "Something went wrong while installing InvokeAI and/or its requirements."
	echo "Update cannot continue. Please report this error to https://github.com/invoke-ai/InvokeAI/issues"
        echo -e "Error code $1; Error caught was '$2'"
        read -p "Press any key to exit..."
        exit
    fi
}

if ! curl -I "$INVOKE_AI_DEP" -fs >/dev/null; then
    echo \'$INVOKE_AI_VERSION\' is not a known branch name or tag. Please check the version and try again.
    exit
fi

echo This script will update InvokeAI and all its dependencies to version \'$INVOKE_AI_VERSION\'.
echo If you do not want to do this, press control-C now!
read -p "Press any key to continue, or CTRL-C to exit..."

curl -L "$INVOKE_AI_DEP" > environments-and-requirements/requirements-base.txt
curl -L "$INVOKE_AI_MODELS" > configs/INITIAL_MODELS.yaml

. .venv/bin/activate

./.venv/bin/python -mpip install -r requirements.txt
_err_exit $? "The pip program failed to install InvokeAI's requirements."

./.venv/bin/python -mpip install $INVOKE_AI_SRC
_err_exit $? "The pip program failed to install InvokeAI."

echo InvokeAI updated to \'$INVOKE_AI_VERSION\'