#!/bin/bash

function show_help() {
    echo "
Usage: install.sh [options]
   Install the TripleBlind SDK.  This supports Linux and MacOS.

Options:
    -sudo          Permits the script to run as root (call sudo ./install.sh -sudo).
    -ni [token]    Non-interactive mode, auto answers Yes to all questions.
                   Optionally include an SDK token to update tripleblind.yaml
    -h, --help     Show this message.
"
}

TOKEN=""

for var in "$@" ; do
    if [[ $var == '-h' || $var == '--help' || $var == '-?' ]] ; then
        show_help
        exit 0
    elif [[ $var == '-sudo' ]] ; then
        sudo_flag="Y"
    elif [[ $var == '-ni' ]] ; then
        non_interactive="Y"
    else
        if [[ $non_interactive == "Y" && $TOKEN == "" ]] ; then
            TOKEN=$var
        else
            echo "Unknown option: $var"
            show_help
            exit 1
        fi
    fi
done


################################################################################
# General script setup

set -e
cd "$(dirname "$0")"
TOP=$(pwd -L)

# exit if run as sudo unless overriden
if [[ "$EUID" -eq 0 && "$sudo_flag" != "Y" ]] ; then
  echo "You are running the installation script as root."
  echo "this will install conda as root, which might be unintuitive or difficult to use."
  echo "If this is still desired, please re-run the script with the -sudo flag."
  exit
fi

if [[ $(sysctl -n machdep.cpu.brand_string 2>/dev/null) =~ "Apple" ]]; then
   IS_M1="Y"
fi

function get_YN() {
    if [[ $non_interactive == "Y" ]] ; then
        return 0  # Answer "Yes" automatically
    fi

    if [[ -z "$1" ]] ; then
        prompt="Choice"
    else
        prompt="$1"
    fi

    if [[ -z "$2" ]] ; then
        yes_msg="Yes"
    else
        yes_msg="$2"
    fi

    if [[ -z "$3" ]] ; then
        no_msg="No"
    else
        no_msg="$3"
    fi

    # Loop until the user hits the Y or the N key
    printf "${prompt} [${CYAN}Y${RESET}/${CYAN}N${RESET}]? "
    while true; do
        if [[ "$OSTYPE" == "darwin"* ]] ; then
            read key
        else
            read -N1 -s key
        fi
        case $key in
        [Yy])
            echo "${HIGHLIGHT}$key - ${yes_msg}${RESET}"
            return 0
            ;;
        [Nn])
            echo "${HIGHLIGHT}$key - ${no_msg}${RESET}"
            return 1
            ;;
        esac
        if [[ "$OSTYPE" == "darwin"* ]] ; then
            printf "${prompt} [${CYAN}Y${RESET}/${CYAN}N${RESET}]: "
        fi
    done
}

function prompt() {
    if [[ $non_interactive == "Y" ]] ; then
        return 0
    fi
    echo "$1" | fmt -w $WIDTH
}

function found_exe() {
    hash "$1" 2>/dev/null
}

function require_rosetta() {
    if ! /usr/bin/pgrep oahd >/dev/null 2>&1; then  # look for Rosetta
        echo "Installing Rosetta for needed x86 support"
        /usr/sbin/softwareupdate --install-rosetta #--agree-to-license
        if [[ $? -ne 0 ]]; then
            echo "Rosetta installation failed, Conda cannot be installed without it."
            exit 11
        fi
    fi
}

if found_exe tput ; then
    if [[ $(tput colors) != "-1" ]]; then
        # Get some colors we can use to spice up messages!
        GREEN=$(tput setaf 2)
        BLUE=$(tput setaf 4)
        BLUE2=$(tput setaf 57)
        BLUE3=$(tput setaf 21)
        RED=$(tput setaf 1)
        CYAN=$(tput setaf 6)
        YELLOW=$(tput setaf 3)
        BACKGROUND=$(tput setab 15)
        RESET=$(tput sgr0)
        HIGHLIGHT=$YELLOW
        PROGRESS=$BLUE
        SUCCESS=$GREEN
    fi

    WIDTH=$(tput cols)
    if (( WIDTH > 100 )) ; then
        WIDTH=100
    fi
else
    WIDTH=78
fi

################################################################################
# Start of installation
echo "${BLUE}${BACKGROUND}"
echo "                  "
echo "        ##        "
echo "       ####       "
echo "        ##        "
echo "                      ${BLUE3}TripleBlind SDK Installer for Linux/MacOS"
echo "   ${BLUE2}##        ${BLUE3}##   "
echo "  ${BLUE2}####      ${BLUE3}####  "
echo "   ${BLUE2}##        ${BLUE3}##   "
echo "                  ${RESET}"

if [[ "$OSTYPE" == "msys" || "$OSTYPE" == "cygwin" ]] ; then
    echo "${YELLOW}WARNING: Unexpected operating system.${RESET}"
    prompt "You are running under Windows, but this script was designed
for a machine running the MacOS or Linux operating system. Windows
users should install via the ${BLUE}install.ps1${RESET} script instead.
"
    if ! get_YN "Do you still want to continue running this script (not recommended)"; then
        exit 1
    fi
fi


prompt "
This script helps you safely install the TripleBlind development environment
on your machine.  It will interactively walk you through prompts to confirm
changes before they are made.  To run without prompts, cancel and rerun with the
'-ni' (non-interactive) flag.

We strongly recommend using a virtual environment to avoid conflicts between
Python versions and libraries. The Conda tool combines a virtual environment
with a package manager, eliminating the pain of library and version conflicts.

This setup will verify the existence of Anaconda or install the minimal
Miniconda bootstrap version if not found, then it will create a virtual
environment for working with TripleBlind's library.
"
if ! get_YN "Would you like to continue the setup"; then
    exit 1
fi

echo
if found_exe conda ; then
    echo "Found Conda version: $(conda --version)"
    eval "$(conda shell.bash hook)"
else
    prompt "Conda was not found on this system.  Miniconda can be downloaded
and installed automatically.  It will take around 250 MB of disk space."

    if ! get_YN "May we download and install miniconda?"; then
        exit 1
    fi

    if [[ $IS_M1 == "Y" ]] ; then
        # We can't run natively on M1s, too many issues still. Require rosetta.
        require_rosetta
    fi

    # create a temporary directory for any transient files
    tmp_d=$(mktemp -d)

    echo "${PROGRESS}Downloading Anaconda distribution...${RESET}"
    tmp_f="$tmp_d/anaconda.sh"
    if [[ $IS_M1 == "Y" ]] ; then
        curl -k https://repo.anaconda.com/miniconda/Miniconda3-py39_4.12.0-MacOSX-arm64.sh -o "$tmp_f"
    elif [[ "$OSTYPE" == "linux-gnu"* ]]; then
        wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-py39_4.12.0-Linux-x86_64.sh -O "$tmp_f" --show-progress
    elif [[ "$OSTYPE" == "darwin"* ]]; then
        curl -k https://repo.anaconda.com/miniconda/Miniconda3-py39_4.12.0-MacOSX-x86_64.sh -o "$tmp_f"
    fi
    chmod +x "$tmp_f"

    # Preferred location is /opt/anaconda3 but the installer cannot install there on MacOS
    # and on Linux this script would need to figure out the proper permissions.
    # For simplicity just installing in home directory.
    if [[ -d ~/anaconda3 ]] ; then
        echo "${PROGRESS}Upgrading Anaconda...${RESET}"
        "$tmp_f" -b -u -p ~/anaconda3
    else
        echo "${PROGRESS}Installing Anaconda...${RESET}"
        "$tmp_f" -b -p ~/anaconda3
    fi

    rm -r "$tmp_d"
    echo "${SUCCESS}Anaconda installation complete!${RESET}"

    echo "${PROGRESS}Adding conda init to shell dotfiles and reloading so conda is available...${RESET}"
    # detect if in bash and init conda for bash
    if [[ "$SHELL" == "/bin/bash" ]] ; then
        ~/anaconda3/bin/conda init bash -q
    fi
    # detect if in zsh and init conda for zsh
    if [[ "$SHELL" == "/bin/zsh" ]] ; then
        ~/anaconda3/bin/conda init zsh -q
    fi

    # Allow conda to satisfy dependency constraints with pip packages
    ~/anaconda3/bin/conda config --set pip_interop_enabled True
    eval "$(~/anaconda3/bin/conda shell.bash hook)"

    echo "${SUCCESS}Miniconda successfully installed!${RESET}"
fi

#################################################
# Check for existing 'tripleblind' environment
if conda env list | grep -q '^tripleblind[[:space:]]'; then
    # Found
    echo
    echo "An existing 'tripleblind' environment was found.  Starting from a"
    echo "clean environment is generally safest.  However, this will remove"
    echo "any extra packages you might have installed manually (pip or conda)."
    echo "You will need to reinstall any such packages after this update."
    echo
    if get_YN "Would you like to start from a clean 'tripleblind' environment?"; then
        conda env remove --name tripleblind
        CREATE_ENV="Y"
    fi
else
    CREATE_ENV="Y"
fi

#################################################

# Activate base for all further commands
conda activate base

# Create a new env if needed
if [[ $CREATE_ENV == "Y" ]] ; then
    conda create --name tripleblind python=3.10 -y
fi

# Update environment
if [[ $IS_M1 == "Y" ]]; then
    # Example: remove or modify xgboost if needed for M1
    cat environment.yaml | grep -v "xgboost" > environment-m1.yaml
    conda env update --name tripleblind -f environment-m1.yaml
    rm environment-m1.yaml
else
    conda env update --name tripleblind -f environment.yaml
fi

#################################################
# Check for existing 'tripleblind' environment script
if [[ ! -f ~/.profile_tripleblind ]] ; then
    echo
    echo "We will create a ~/.profile_tripleblind which has commands to make"
    echo "working with TripleBlind easy.  It can be automatically invoked when"
    echo "you start your terminal session, otherwise you will need to manually"
    echo "source it in order to prepare your environment for working with the SDK."
    echo
    if get_YN "Would you like the profile_tripleblind automatically loaded?"; then

        echo 'source ~/.profile_tripleblind' >> ~/.zshrc
        if [[ "$OSTYPE" == "darwin"* ]] ; then
            # On a Mac, the .bash_profile is loaded when you start the terminal
            echo 'source ~/.profile_tripleblind' >> ~/.bash_profile
        else
            # Linux loads .bashrc
            echo 'source ~/.profile_tripleblind' >> ~/.bashrc
        fi
    else
        echo "To work with the SDK, you will need to run the command:"
        echo "   source ~/.profile_tripleblind"
    fi
fi

# Create/overwrite the ~/.profile_tripleblind file
cat <<EOF > ~/.profile_tripleblind
# WARNING: This file may be replaced in future, do not customize.

export TB_API_CONFIG='${TOP}/tripleblind.yaml'
alias tb='python ${TOP}/tb.py'
if [ "\$PS1" ] ; then  # only execute within interactive shells
    conda activate tripleblind
fi
EOF

echo "${SUCCESS}Profile created!${RESET}"

#################################################
# Setup the user access token

echo
echo "In order to associate the SDK with your user account, you must use"
echo "your personalized tripleblind.yaml file."
echo
echo "To collect your tripleblind.yaml file:"
echo "  - Browse to http://tripleblind.app and log in"
echo "  - Click on the profile icon in the upper-right corner and select My Settings"
echo "  - In the Authentication tab, click Download tripleblind.yaml"
echo "  - Place the downloaded file into your SDK folder (replace the existing file)"
echo "Alternatively, paste just your Authentication Token below."

if [[ $non_interactive != "Y" ]] ; then
    current_token=$(grep 'sdkuser_token: ".*"' tripleblind.yaml | grep -o '".*"' || true)

    if [[ "$current_token" == '""' ]] || [[ -z "$current_token" ]] ; then
        echo -n "If you have set up your yaml file, press Enter. Or, provide your token: "
        read TOKEN
    else
        echo "Current access token: $current_token"
        echo -n "Enter a new access token or blank to keep: "
        read TOKEN
    fi
fi

if [[ "$TOKEN" != "" ]]; then
    sed "s/sdkuser_token: \".*\"/sdkuser_token: \"$TOKEN\"/" < tripleblind.yaml > tripleblind.tmp
    mv tripleblind.tmp tripleblind.yaml
    echo "${SUCCESS}Token set!  See tripleblind.yaml for more details.${RESET}"
fi

echo
echo "${PROGRESS}"
echo "***********************************************************************"
echo "Installation successful!  You are ready to use the TripleBlind SDK.  Be"
echo "sure to use the 'tripleblind' conda environment whenever working with it."
echo "    conda activate tripleblind"
echo
echo "${HIGHLIGHT}IMPORTANT:${PROGRESS} You need to close and restart your terminal once before you"
echo "can use the 'conda' command."
echo
echo "Open the README.html for documentation and help getting started.${RESET}"
