#!/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=$@

# This setting allows torch to fall back to CPU for operations that are not supported by MPS on macOS.
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 "Open the developer console\n"
        file_name=$(basename "${BASH_SOURCE[0]}")
        bash --init-file "$file_name"
        ;;
    3)
        clear
        printf "Running the db maintenance script\n"
        invokeai-db-maintenance --root ${INVOKEAI_ROOT}
        ;;
    4)
        clear
        printf "Command-line help\n"
        invokeai-web --help
        ;;
    *)
        clear
        printf "Exiting...\n"
        exit
        ;;
    esac
    clear
}

# Command-line interface for launching Invoke functions
do_line_input() {
    clear
    printf "What would you like to do?\n"
    printf "1: Generate images using the browser-based interface\n"
    printf "2: Open the developer console\n"
    printf "3: Run the InvokeAI image database maintenance script\n"
    printf "4: Command-line help\n"
    printf "Q: Quit\n\n"
    printf "To update, download and run the installer from https://github.com/invoke-ai/InvokeAI/releases/latest.\n\n"
    read -p "Please enter 1-4, Q: [1] " yn
    choice=${yn:='1'}
    do_choice $choice
    clear
}

# Main IF statement for launching Invoke, and for checking if the user is in the developer console
if [ "$0" != "bash" ]; then
    while true; do
        do_line_input
    done
else # in developer console
    python --version
    printf "Press ^D to exit\n"
    export PS1="(InvokeAI) \u@\h \w> "
fi