#!/bin/bash

# MIT License

# Coauthored by Lincoln Stein, Eugene Brodsky and Joshua Kimsey
# 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=$@

# Check to see if dialog is installed (it seems to be fairly standard, but good to check regardless) and if the user has passed the --no-tui argument to disable the dialog TUI
tui=true
if command -v dialog &>/dev/null; then
    # This must use $@ to properly loop through the arguments passed by the user
    for arg in "$@"; do
        if [ "$arg" == "--no-tui" ]; then
            tui=false
            # Remove the --no-tui argument to avoid errors later on when passing arguments to InvokeAI
            PARAMS=$(echo "$PARAMS" | sed 's/--no-tui//')
            break
        fi
    done
else
    tui=false
fi

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

# Avoid glibc memory fragmentation. See invokeai/backend/model_management/README.md for details.
export MALLOC_MMAP_THRESHOLD_=1048576

# Primary function for the case statement to determine user input
do_choice() {
    case $1 in
    1)
        clear
        printf "Generate images with a browser-based interface\n"
        invokeai-web $PARAMS
        ;;
    2)
        clear
        printf "Explore InvokeAI nodes using a command-line interface\n"
        invokeai $PARAMS
        ;;
    3)
        clear
        printf "Textual inversion training\n"
        invokeai-ti --gui $PARAMS
        ;;
    4)
        clear
        printf "Merge models (diffusers type only)\n"
        invokeai-merge --gui $PARAMS
        ;;
    5)
        clear
        printf "Download and install models\n"
        invokeai-model-install --root ${INVOKEAI_ROOT}
        ;;
    6)
        clear
        printf "Change InvokeAI startup options\n"
        invokeai-configure --root ${INVOKEAI_ROOT} --skip-sd-weights --skip-support-models
        ;;
    7)
        clear
        printf "Re-run the configure script to fix a broken install or to complete a major upgrade\n"
        invokeai-configure --root ${INVOKEAI_ROOT} --yes --default_only --skip-sd-weights
        ;;
    8)
        clear
        printf "Open the developer console\n"
        file_name=$(basename "${BASH_SOURCE[0]}")
        bash --init-file "$file_name"
        ;;
    9)
        clear
        printf "Update InvokeAI\n"
        python -m invokeai.frontend.install.invokeai_update
        ;;
    10)
        clear
        printf "Running the db maintenance script\n"
        invokeai-db-maintenance --root ${INVOKEAI_ROOT}
        ;;
    11)
        clear
        printf "Command-line help\n"
        invokeai-web --help
        ;;
    *)
        clear
        printf "Exiting...\n"
        exit
        ;;
    esac
    clear
}

# Dialog-based TUI for launcing Invoke functions
do_dialog() {
    options=(
        1 "Generate images with a browser-based interface"
        2 "Explore InvokeAI nodes 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 or to complete a major upgrade"
        8 "Open the developer console"
        9 "Update InvokeAI"
	10 "Run the InvokeAI image database maintenance script"
	11 "Command-line help"
    )

    choice=$(dialog --clear \
        --backtitle "\Zb\Zu\Z3InvokeAI" \
        --colors \
        --title "What would you like to do?" \
        --ok-label "Run" \
        --cancel-label "Exit" \
        --help-button \
        --help-label "CLI Help" \
        --menu "Select an option:" \
        0 0 0 \
        "${options[@]}" \
        2>&1 >/dev/tty) || clear
    do_choice "$choice"
    clear
}

# Command-line interface for launching Invoke functions
do_line_input() {
    clear
    printf " ** For a more attractive experience, please install the 'dialog' utility using your package manager. **\n\n"
    printf "What would you like to do?\n"
    printf "1: Generate images using the browser-based interface\n"
    printf "2: Explore InvokeAI nodes using the command-line interface\n"
    printf "3: Run textual inversion training\n"
    printf "4: Merge models (diffusers type only)\n"
    printf "5: Download and install models\n"
    printf "6: Change InvokeAI startup options\n"
    printf "7: Re-run the configure script to fix a broken install\n"
    printf "8: Open the developer console\n"
    printf "9: Update InvokeAI\n"
    printf "10: Run the InvokeAI image database maintenance script\n"
    printf "11: Command-line help\n"
    printf "Q: Quit\n\n"
    read -p "Please enter 1-11, Q: [1] " yn
    choice=${yn:='1'}
    do_choice $choice
    clear
}

# Main IF statement for launching Invoke with either the TUI or CLI, and for checking if the user is in the developer console
if [ "$0" != "bash" ]; then
    while true; do
        if $tui; then
            # .dialogrc must be located in the same directory as the invoke.sh script
            export DIALOGRC="./.dialogrc"
            do_dialog
        else
            do_line_input
        fi
    done
else # in developer console
    python --version
    printf "Press ^D to exit\n"
    export PS1="(InvokeAI) \u@\h \w> "
fi