#!/bin/bash

# coauthored by Lincoln Stein, Eugene Brodsky and JoshuaKimsey
# Copyright 2023, The InvokeAI Development Team

####
# This launch script assumes that:
# 1. it is located in the runtime directory,
# 2. the .venv is also located in the runtime directory and is named exactly that
#
# If both of the above are not true, this script will likely not work as intended.
# Activate the virtual environment and run `invoke.py` directly.
####

set -eu

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

. .venv/bin/activate

export INVOKEAI_ROOT="$scriptdir"
PARAMS=$@

# set required env var for torch on mac MPS
if [ "$(uname -s)" == "Darwin" ]; then
    export PYTORCH_ENABLE_MPS_FALLBACK=1
fi

do_choice() {
    case $1 in
	1)
            echo "Generate images with a browser-based interface"
            clear
            invokeai --web $PARAMS
            ;;
	2)
            echo "Generate images using a command-line interface"
            clear
            invokeai $PARAMS
            ;;
	3)
            echo "Textual inversion training"
            clear
            invokeai-ti --gui $PARAMS
            ;;
	4)
            echo "Merge models (diffusers type only)"
            clear
            invokeai-merge --gui $PARAMS
            ;;
	5)
            echo "Download and install models"
            clear
            invokeai-model-install --root ${INVOKEAI_ROOT}
            ;;
	6)
            echo "Change InvokeAI startup options"
            clear
            invokeai-configure --root ${INVOKEAI_ROOT} --skip-sd-weights --skip-support-models
            ;;
	7)
            echo "Re-run the configure script to fix a broken install"
            clear
            invokeai-configure --root ${INVOKEAI_ROOT} --yes --default_only
            ;;
	8)
            echo "Open the developer console"
            clear
            file_name=$(basename "${BASH_SOURCE[0]}")
            bash --init-file "$file_name"
            ;;
	9)
            echo "Update InvokeAI"
            clear
            invokeai-update
            ;;
	10)
            echo "Command-line help"
            clear
            invokeai --help
            ;;
	*)
            echo "Exiting..."
            exit
            ;;
    esac
    clear
}
    
do_dialog() {
    while true
    do
	options=(
	    1 "Generate images with a browser-based interface"
	    2 "Generate images using a command-line interface"
	    3 "Textual inversion training"
	    4 "Merge models (diffusers type only)"
	    5 "Download and install models"
	    6 "Change InvokeAI startup options"
	    7 "Re-run the configure script to fix a broken install"
	    8 "Open the developer console"
	    9 "Update InvokeAI"
	    10 "Command-line help")
	
        choice=$(dialog --clear \
                        --backtitle "InvokeAI" \
                        --title "What would you like to run?" \
                        --menu "Select an option:" \
                        0 0 0 \
                        "${options[@]}" \
                        2>&1 >/dev/tty) || clear
	do_choice "$choice"
    done
    clear
}

do_line_input() {
    echo " ** For a more attractive experience, please install the 'dialog' utility. **"
    echo ""
    while true
    do
	echo "Do you want to generate images using the"
	echo "1. browser-based UI"
	echo "2. command-line interface"
	echo "3. run textual inversion training"
	echo "4. merge models (diffusers type only)"
	echo "5. download and install models"
	echo "6. change InvokeAI startup options"
	echo "7. re-run the configure script to fix a broken install"
	echo "8. open the developer console"
	echo "9. update InvokeAI"
	echo "10. command-line help"
	echo "Q - Quit"
	echo ""
	read -p "Please enter 1-10, Q: [1] " yn
	choice=${yn:='1'}
	do_choice $choice
    done
}

if [ "$0" != "bash" ]; then
    # Dialog seems to be a standard installtion for most Linux distros, but this checks to ensure it is present regardless
    if command -v dialog &> /dev/null ; then
	do_dialog
    else
	do_line_input
    fi
else # in developer console
    python --version
    echo "Press ^D to exit"
    export PS1="(InvokeAI) \u@\h \w> "
fi