#!/bin/bash
# This file is part of the jetson_stats package (https://github.com/rbonghi/jetson_stats or http://rnext.it).
# Copyright (c) 2019 Raffaello Bonghi.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.

red=`tput setaf 1`
green=`tput setaf 2`
reset=`tput sgr0`

# Load sudo user nane
if [ ! -z $SUDO_USER ] ; then
    USER=$SUDO_USER
fi
TITLE="jetson-stats configurator"
INTERACTIVE=true
FORCE=false
# List of all health functions
# For each health name are required three differents functions:
# - status_<name>     -> Check the status of health function
# - fix_<name>        -> Run the installer
# - uninstall_<name>  -> Run the uninstaller
# - <name>            -> Make the gui
HEALTH_FUNCTIONS=("commands" "Check all commands if are installed"\
                  "folders" "Check folders"\
                  "variables" "Check variables status"\
                  "permissions" "Check status permissions for $USER"\
                  "service" "Status jetson-stats service")
# Jetson-stats main folder
JETSON_STATS_FOLDER="/usr/local/jetson_stats"
# Locate jetson-stats and python version
# Load JETSON environment variables
JTOP_VARIABLE=""
JTOP_PYTHON=""
JTOP_PYTHON_NAME=""
if type -P python3 >/dev/null 2>&1 ; then
    JTOP_VARIABLE=$(python3 -c "import jtop; print(jtop.__path__[0])" 2> /dev/null)
    JTOP_PYTHON="$(python3 -V 2>&1)"
    JTOP_PYTHON_NAME="python3"
fi
if type -P python >/dev/null 2>&1 && [ -z $JTOP_VARIABLE ] ; then
    JTOP_VARIABLE=$(python -c "import jtop; print(jtop.__path__[0])" 2> /dev/null)
    JTOP_PYTHON="$(python -V 2>&1)"
    JTOP_PYTHON_NAME="python"
fi
# Load jetson_variables
if [ ! -z $JTOP_VARIABLE ] ; then
    source $JTOP_VARIABLE/jetson_variables
fi

status_commands()
{
    local BIN_FOLDER="/usr/local/bin"
    if [ ! -f "$BIN_FOLDER/jetson_release" ] ; then
        echo 1
        return
    fi
    if [ ! -f "$BIN_FOLDER/jetson_swap" ] ; then
        echo 1
        return
    fi
    # Everything is fine
    echo 0
}

commands()
{
    local TEXT=""
    # Check status folders
    if [ $(status_commands) -eq 1 ] ; then
        TEXT="Require to reinstall this package using\n"
        TEXT+="sudo -H pip install --no-cache-dir -U jetson-stats"
        # Run script to fix variables
        if (whiptail --title "$TITLE" --yesno "$TEXT\n\nDo you want fix?" 11 78); then
            fix_commands
        fi
    else
        TEXT="All commands are correctly installed"
        # Run Message
        whiptail --title "$TITLE" --msgbox "$TEXT" 10 78
    fi
    # Return to healt function
    health
}

fix_commands()
{
    # Skip this script if forced
    if $FORCE ; then
        echo "${red}Skip with FORCE option${reset}"
        return
    fi
    # Recognize version pip
    local PIP_VERSION="pip"
    if [ $JTOP_PYTHON_NAME == "python3" ] ; then
        PIP_VERSION="pip3"
    fi
    # Uninstall package
    sudo -H $PIP_VERSION uninstall -y jetson-stats
    # Reinstall package without cache dir pip
    sudo -H $PIP_VERSION install --no-cache-dir -U jetson-stats
    # After this script will be close the script
    exit 0
}

uninstall_commands()
{
    echo "Nothing to do"
}

status_folders()
{
    local BIN_FOLDER="/usr/local/bin"
    # Old jetson_names
    if [ -f "$BIN_FOLDER/jetson-docker" ] || [ -f "$BIN_FOLDER/jetson-release" ] || [ -f "$BIN_FOLDER/jetson-swap" ] ; then
        echo 1
        return
    fi
    # Old folder jetson_stats
    if [ -d "/etc/jetson_easy" ] ; then
        echo 1
        return
    fi
    if [ -d "/opt/jetson_stats" ] ; then
        echo 1
        return
    fi
    if [ ! -d $JETSON_STATS_FOLDER ] ; then
        echo 1
        return
    fi
    # Everything is fine
    echo 0
}

fix_folders()
{
    local BIN_FOLDER="/usr/local/bin"
    # Remove jetson-docker link
    if [ -f "$BIN_FOLDER/jetson-docker" ] ; then
        sudo rm "$BIN_FOLDER/jetson-docker"
    fi
    # Remove jetson_release link
    if [ -f "$BIN_FOLDER/jetson-release" ] ; then
        sudo rm "$BIN_FOLDER/jetson-release"
    fi
    # Remove jetson_release link
    if [ -f "$BIN_FOLDER/jetson-swap" ] ; then
        sudo rm "$BIN_FOLDER/jetson-swap"
    fi
    # Remove /etc/jetson_easy folder
    if [ -d "/etc/jetson_easy" ] ; then
        sudo rm -R "/etc/jetson_easy"
    fi
    # Remove jetson_release link
    if [ -d "/etc/jetson_easy" ] ; then
        sudo rm -R "/etc/jetson_easy"
    fi
    # Remove jetson_release link
    if [ -d "/opt/jetson_stats" ] ; then
        # remove old reference folders
        if [ -f "$BIN_FOLDER/jetson_docker" ] ; then
            sudo rm "$BIN_FOLDER/jetson_docker"
        fi
        # Remove jetson_release link
        if [ -f "$BIN_FOLDER/jetson_release" ] ; then
            sudo rm "$BIN_FOLDER/jetson_release"
        fi
        # Remove jetson_release link
        if [ -f "$BIN_FOLDER/jetson_swap" ] ; then
            sudo rm "$BIN_FOLDER/jetson_swap"
        fi
        # Remove reference folder
        sudo rm -R "/opt/jetson_stats"
    fi

}

uninstall_folders()
{
    echo "Nothing to do"
}

folders()
{
    local TEXT=""
    # Check status folders
    if [ $(status_folders) -eq 1 ] ; then
        TEXT="Clean and reinstall all folders."
        # Run script to fix variables
        if (whiptail --title "$TITLE" --yesno "$TEXT\n\nDo you want fix?" 11 78); then
            fix_folders
        fi
    else
        TEXT="All folder are correctly installed"
        # Run Message
        whiptail --title "$TITLE" --msgbox "$TEXT" 10 78
    fi
    # Return to healt function
    health
}

fix_permissions()
{
    local group="jetson_stats"
    # Make jetson_stats group
    sudo groupadd $group
    # Add jetson_stats user group to local user
    sudo usermod -a -G $group $USER
    # Require reboot
    if [ ! -f /var/run/reboot-required ] ; then
        sudo sh -c 'echo "*** System Restart Required ***" > /var/run/reboot-required'
    fi
}

uninstall_permissions()
{
    echo "Nothing to do"
}

status_permissions()
{
    local group="jetson_stats"
    # Check if exist group
    if ! grep -q $group /etc/group ; then
        echo 1
        return
    fi
    # Check if user is in group
    if ! groups $USER | grep -q '\bjetson_stats\b'; then
        echo 1
        return
    fi
    # Everything is fine
    echo 0
}

permissions()
{
    # Load text
    local TEXT="jtop to read all data from jetson_stats service \
need to be in jetson_stats group."
    # Check status status_variables
    if [ $(status_permissions) -eq 1 ] ; then
        TEXT+="\n\nThe does not have proper permissions.\nDo you want fix?"
        # Run script to fix variables
        if (whiptail --title "$TITLE" --yesno "$TEXT" 11 78); then
            fix_service
        fi
    else
        # Run Message
        whiptail --title "$TITLE" --msgbox "$TEXT" 10 78
    fi
    # Return to healt function
    health
}

fix_service()
{
    # Remove old jetson_performance.service and jetson_stats_boot.service
    if [ -f "/etc/systemd/system/jetson_performance.service" ] ; then
        # Uninstall the service
        if [ $(systemctl is-active jetson_performance.service) = "active" ] ; then
            # Stop the service
            sudo systemctl stop jetson_performance.service
        fi
        # Disable the service
        sudo systemctl disable jetson_performance.service
        # Remove jetson_performance service from /etc/init.d
        sudo rm "/etc/systemd/system/jetson_performance.service"
        # Update service list
        sudo systemctl daemon-reload
    fi
    # Disable service
    if [ -f "/etc/systemd/system/jetson_stats_boot.service" ] ; then
        sudo systemctl disable jetson_stats_boot.service
        # Remove fan service from /etc/init.d
        sudo rm "/etc/systemd/system/jetson_stats_boot.service"
        # Update service list
        sudo systemctl daemon-reload
    fi
    # Using a copy of the service to prevent the remotion when is call the function "disable"
    sudo cp "$JETSON_STATS_FOLDER/jetson_stats.service" "/etc/systemd/system/jetson_stats.service"
    # Old configuration
    # sudo ln -s "$JETSON_STATS_FOLDER/jetson_stats.service" "/etc/systemd/system/jetson_stats.service" 2> /dev/null
    # Update service list
    sudo systemctl daemon-reload
    # Enable jetson_stats at startup
    sudo systemctl enable jetson_stats.service
    # Start service
    sudo systemctl start jetson_stats.service
}

uninstall_service()
{
    # Remove old jetson_performance.service and jetson_stats_boot.service
    if [ -f "/etc/systemd/system/jetson_performance.service" ] ; then
        # Uninstall the service
        if [ $(systemctl is-active jetson_performance.service) = "active" ] ; then
            # Stop the service
            sudo systemctl stop jetson_performance.service
            echo "Stop jetson_performance.service"
        fi
        # Disable the service
        sudo systemctl disable jetson_performance.service
        # Remove jetson_performance service from /etc/init.d
        sudo rm "/etc/systemd/system/jetson_performance.service"
        # Update service list
        sudo systemctl daemon-reload
    fi
    # Disable old service jetson_stats_boot
    if [ -f "/etc/systemd/system/jetson_stats_boot.service" ] ; then
        sudo systemctl disable jetson_stats_boot.service
        # Remove fan service from /etc/init.d
        sudo rm "/etc/systemd/system/jetson_stats_boot.service"
        # Update service list
        sudo systemctl daemon-reload
    fi
    # Disable service jetson_stats.service
    if [ -f "/etc/systemd/system/jetson_stats.service" ] ; then
        # Uninstall the service
        if [ $(systemctl is-active jetson_stats.service) = "active" ] ; then
            # Stop the service
            sudo systemctl stop jetson_stats.service
            echo "Stop jetson_stats.service"
        fi
        # Disable the service
        sudo systemctl disable jetson_stats.service
        # Remove fan service from /etc/init.d
        sudo rm "/etc/systemd/system/jetson_stats.service"
        # Update service list
        sudo systemctl daemon-reload
    fi
    # Remove all jtop files
    if [ -f "/run/jtop.sock" ] ; then
        sudo rm "/run/jtop.sock"
    fi
}

status_service()
{
    # Check if old services are installed
    if [ -f "/etc/systemd/system/jetson_performance.service" ] ; then
        echo 1
        return
    fi
    if [ -f "/etc/systemd/system/jetson_stats_boot.service" ] ; then
        echo 1
        return
    fi
    # Check if is installed jetson_stats.service
    if [ ! -f "/etc/systemd/system/jetson_stats.service" ] ; then
        echo 1
        return
    fi
    # Check if file are same
    if ! cmp --silent "/etc/systemd/system/jetson_stats.service" "$JETSON_STATS_FOLDER/jetson_stats.service" ; then
        echo 1
        return
    fi
    # Check status service
    local JETSON_STATS_STATUS="$(systemctl is-active jetson_stats.service)"
    if [ $JETSON_STATS_STATUS != "active" ] ; then
        echo 1
        return
    fi
    local JETSON_STATS_ENABLED=$(systemctl is-enabled jetson_stats.service 2> /dev/null)
    if [ -z $JETSON_STATS_ENABLED ] ; then
        JETSON_STATS_ENABLED="failed"
        echo 1
        return
    fi
    # Everything is fine
    echo 0
}

service()
{
    local TEXT="jtop use a service to interact with your NVIDIA jetson \
When you install jtop you can find it in systemctl jetson_stats.service"
    # Check status status_variables
    if [ $(status_service) -eq 1 ] ; then
        TEXT+="\n\nThe jetson_stats service is in failure\nDo you want fix?"
        # Run script to fix variables
        if (whiptail --title "$TITLE" --yesno "$TEXT" 11 78); then
            fix_service
        fi
    else
        # Run Message
        whiptail --title "$TITLE" --msgbox "$TEXT" 10 78
    fi
    # Return to healt function
    health
}

status_variables()
{
    # Check if file does not exist
    if [ ! -f "/etc/profile.d/jetson_env.sh" ] ; then
        echo 1
        return
    fi
    # Check if file are different
    if ! cmp --silent "/etc/profile.d/jetson_env.sh" "$JETSON_STATS_FOLDER/jetson_env.sh" ; then
        echo 1
        return
    fi
    # Everything is fine
    echo 0
}

fix_variables()
{
    sudo cp "$JETSON_STATS_FOLDER/jetson_env.sh" "/etc/profile.d/jetson_env.sh"
    # Require reboot
    if [ ! -f /var/run/reboot-required ] ; then
        sudo sh -c 'echo "*** System Restart Required ***" > /var/run/reboot-required'
    fi
}

uninstall_variables()
{
    # Remove jetson_env variable
    if [ -f "/etc/profile.d/jetson_env.sh" ] ; then
        sudo rm "/etc/profile.d/jetson_env.sh"
        # Require reboot
        if [ ! -f /var/run/reboot-required ] ; then
            sudo sh -c 'echo "*** System Restart Required ***" > /var/run/reboot-required'
        fi
    else
        echo "No jetson environment variables scripts installed"
    fi
}

variables()
{
    local TEXT="jetson-stats add new enviroment variables in your bash starting \
with \"JETSON_\" with detailed information about your board.\nThey are located in \"/etc/profile.d/jetson_env.sh"
    # Check status status_variables
    if [ $(status_variables) -eq 1 ] ; then
        TEXT+="\n\nThe Jetson enviroment variables are in failure\nDo you want fix?"
        # Run script to fix variables
        if (whiptail --title "$TITLE" --yesno "$TEXT" 11 78); then
            fix_variables
        fi
    else
        # Run Message
        whiptail --title "$TITLE" --msgbox "$TEXT" 10 78
    fi
    # Return to healt function
    health
}

installer()
{
    local LIST=("$@")
    # Evaluate length list
    local LIST_LEN
    let LIST_LEN=${#LIST[@]}-1
    # Check if there is a function require fix
    for i in `seq 0 2 $LIST_LEN`; do
        # Read name
        local name=${LIST[$i]}
        local description=${LIST[$i+1]}
        # Evaluate status
        local status=$(eval "status_$name")
        # Uncomment this line only for test
        # echo $i $name $status
        # Run uninstall script
        local status=$(eval "status_$name")
        # Check if a status function in list return error=1
        if [ "$status" -eq 1 ] || $FORCE ; then
            echo "${green}[Install $name - $description]${reset}"
            # Run the Install scripts script
            eval "fix_$name"
        fi
    done
}

uninstaller()
{
    local LIST=("$@")
    # Evaluate length list
    local LIST_LEN
    let LIST_LEN=${#LIST[@]}-1
    # Check if there is a function require fix
    for i in `seq 0 2 $LIST_LEN`; do
        # Read name
        local name=${LIST[$i]}
        local description=${LIST[$i+1]}
        # Evaluate status
        local status=$(eval "status_$name")
        # Uncomment this line only for test
        # echo $i $name $status
        # Run uninstall script
        if type "uninstall_$name"| grep -i function > /dev/null; then
            echo "${green}[Uninstall $name - $description]${reset}"
            # Call Uninstaller function
            eval "uninstall_$name"
        else
            echo "Function uninstall_$name does not exist."
        fi
    done
}

write_status()
{
    if [ $1 -eq 1 ] ; then
        echo "n"
    else
        echo "Y"
    fi
}

status_health()
{
    local LIST=("$@")
    # Evaluate length list
    local LIST_LEN
    let LIST_LEN=${#LIST[@]}-1
    # Check if there is a function require fix
    for i in `seq 0 2 $LIST_LEN`; do
        # Read name
        local name=${LIST[$i]}
        # Evaluate status
        local status=$(eval "status_$name")
        # Check if a status function in list return error=1
        if [ $status -eq 1 ] ; then
            return 1
        fi
    done
    return 0
}

check_run()
{
    local status=$1
    # Check if require ask before exectute
    if [ $status -eq 1 ] ; then
        # If you cannot understand this, read Bash_Shell_Scripting#if_statements again.
        if (whiptail --title "$TITLE" --yesno "This is an example of a yes/no box." 8 78); then
            echo $?
        fi
    fi
}

health()
{
    # Evaluate the size
    local HEALTH_FUNCTIONS_LEN
    let HEALTH_FUNCTIONS_LEN=${#HEALTH_FUNCTIONS[@]}-1
    local OK_BUTTON="Fix"
    # Make menu
    local MENU="jetson-stats configuration\n"
    MENU+=" - Folder scripts: $JETSON_STATS_FOLDER\n"
    MENU+=" - jtop installed on $JTOP_PYTHON\n"
    MENU+=" - jtop path: $JTOP_VARIABLE\n"
    MENU+="Choose an option:"
    # Build Menu healt list
    local MENU_LIST=("<-- Back" "Return to the main menu.")
    # Make jetson_stats menu
    if status_health "${HEALTH_FUNCTIONS[@]}" ; then
        OK_BUTTON="Ok"
    else
        MENU_LIST+=("Fix" "Fix all packages")
    fi
    # Build menu list
    for i in `seq 0 2 $HEALTH_FUNCTIONS_LEN`; do
        # Read name
        local name=${HEALTH_FUNCTIONS[$i]}
        local description=${HEALTH_FUNCTIONS[$i+1]}
        # Evaluate status
        local status=$(eval "status_$name")
        # echo $i $name $status
        MENU_LIST+=("[$(write_status $status)] $name" "$description")
    done
    # Evaluate the size
    local MENU_LIST_LEN
    let MENU_LIST_LEN=${#MENU_LIST[@]}/2
    local OUTPUT
    OUTPUT=$(whiptail --title "$TITLE" --ok-button "$OK_BUTTON" --cancel-button "return" --menu "$MENU" 18 78 $MENU_LIST_LEN "${MENU_LIST[@]}" 3>&1 1>&2 2>&3)
    # Read exit status
    if [ $? -eq 0 ]; then
        # Run the function with the same name
        if [ "$OUTPUT" = "<-- Back" ] ; then
            # Return to main
            main
        elif [ "$OUTPUT" = "Fix" ] ; then
            # Run the installer script
            installer "${HEALTH_FUNCTIONS[@]}"
            # End loop
            do_finish
            # Return to main
            health
        else
            # Extract function name
            local FUNCTION=$(echo $OUTPUT | cut -d " " -f2)
            # Evaluate status
            local status=$(eval "status_$FUNCTION")
            # Run function
            $FUNCTION
            # End loop
            do_finish
        fi
    else
        # Return to main for other options
        main
    fi
}

get_type_desktop()
{
    local GUI_CONF="/etc/gdm3/custom.conf"
    # Check desktop configuration
    local DEFAULT_TARGET=$(systemctl get-default)
    local TYPE=""
    if [ "$DEFAULT_TARGET" = "multi-user.target" ] ; then
        TYPE="B1"
        if [ -f "/etc/systemd/system/getty@tty1.service.d/autologin.conf" ] ; then
            TYPE="B2"
        fi
    elif [ "$DEFAULT_TARGET" = "graphical.target" ] ; then
        TYPE="B3"
        if [ -f $GUI_CONF ] ; then
            if grep -q "AutomaticLoginEnable=[tT]rue" $GUI_CONF ; then
                TYPE="B4"
            fi
        fi
    fi
    echo $TYPE
    return
}

update_desktop_login()
{
    local GUI_CONF="/etc/gdm3/custom.conf"
    local option=$1
    local user=$2
    if [ -f $GUI_CONF ] ; then
        # Remove old Automatic configuration data
        if grep -q "AutomaticLogin=" $GUI_CONF ; then
            sudo sed -i "/^AutomaticLogin=/d" $GUI_CONF
            echo "Removed"
        fi
        # Add new status
        sudo sed -i "/\[daemon\]/a AutomaticLogin=$user" $GUI_CONF

        # Remove old Automatic configuration data
        if grep -q "AutomaticLoginEnable=" $GUI_CONF ; then
            sudo sed -i "/^AutomaticLoginEnable=/d" $GUI_CONF
        fi
        # Add new status
        sudo sed -i "/\[daemon\]/a AutomaticLoginEnable=$option" $GUI_CONF
    else
        echo "${red}$GUI_CONF does not exist!${reset}"
    fi
}

identify_desktop()
{
    local TYPE=$1
    local NAME=$2
    if [ "$TYPE" = "$NAME" ] ; then
        echo "[$NAME]"
    else
        echo "$NAME"
    fi
}

desktop() {
    # Desktop manager. Enable and disable desktop feature with autologin or not
    # Reference:
    # 1. https://github.com/RPi-Distro/raspi-config/blob/master/raspi-config#L1122
    # 2. https://linuxconfig.org/how-to-enable-automatic-login-on-ubuntu-18-04-bionic-beaver-linux
    # 3. https://askubuntu.com/questions/819117/how-can-i-get-autologin-at-startup-working-on-ubuntu-server-16-04-1
    local MENU="Status board:\n"
    local DEFAULT_TARGET=$(systemctl get-default)
    local AUTO_LOGIN="False"
    local REQUIRE_REBOOT=false
    if [ $DEFAULT_TARGET = "multi-user.target" ] ; then
        if [ -f "/etc/systemd/system/getty@tty1.service.d/autologin.conf" ] ; then
            AUTO_LOGIN="True"
        fi
    else
        if [ -f "/etc/gdm3/custom.conf" ] ; then
            if grep -q "AutomaticLoginEnable=[tT]rue" "/etc/gdm3/custom.conf" ; then
                AUTO_LOGIN="True"
            fi
        fi
    fi
    MENU+=" - $DEFAULT_TARGET\n"
    MENU+=" - Autologin: $AUTO_LOGIN"
    local DESKTOP_TYPE=$(get_type_desktop)
    # Make menu menu
    if $INTERACTIVE ; then
        BOOTOPT=$(whiptail --title "$TITLE" --menu "$MENU\n\nBoot Options" 18 78 6 \
            "<-- Back" "Return to the main menu" \
            "$(identify_desktop $DESKTOP_TYPE B1) Console" "Text console, requiring user to login" \
            "$(identify_desktop $DESKTOP_TYPE B2) Console Autologin" "Text console, automatically logged in as '$USER' user" \
            "$(identify_desktop $DESKTOP_TYPE B3) Desktop" "Desktop GUI, requiring user to login" \
            "$(identify_desktop $DESKTOP_TYPE B4) Desktop Autologin" "Desktop GUI, automatically logged in as '$USER' user" \
        3>&1 1>&2 2>&3)
    else
        BOOTOPT=$1
        true
    fi
    # Pass option
    if [ $? -eq 0 ]; then
        case "$BOOTOPT" in
            "<-- Back")
                # Return to main for other options
                ;;
            *B1*)
                sudo systemctl set-default multi-user.target
                sudo ln -fs /lib/systemd/system/getty@.service /etc/systemd/system/getty.target.wants/getty@tty1.service
                # Remove autologin
                if [ -f "/etc/systemd/system/getty@tty1.service.d/autologin.conf" ] ; then
                    sudo rm /etc/systemd/system/getty@tty1.service.d/autologin.conf
                fi
                REQUIRE_REBOOT=true
                ;;
            *B2*)
                sudo systemctl set-default multi-user.target
                sudo ln -fs /lib/systemd/system/getty@.service /etc/systemd/system/getty.target.wants/getty@tty1.service
                # If the file does not exist make it one
                if [ ! -f "/etc/systemd/system/getty@tty1.service.d/autologin.conf" ] ; then
                    sudo mkdir -p "/etc/systemd/system/getty@tty1.service.d"
                    sudo touch "/etc/systemd/system/getty@tty1.service.d/autologin.conf"
                fi
                # Change hownership
                sudo chown $USER:$USER "/etc/systemd/system/getty@tty1.service.d/autologin.conf"
                # Copy the autologin configuration
                cat > /etc/systemd/system/getty@tty1.service.d/autologin.conf << EOF
[Service]
ExecStart=
ExecStart=-/sbin/agetty --autologin $USER --noclear %I \$TERM
EOF
                REQUIRE_REBOOT=true
                ;;
            *B3*)
                sudo systemctl set-default graphical.target
                sudo ln -fs /lib/systemd/system/getty@.service /etc/systemd/system/getty.target.wants/getty@tty1.service
                # Remove autologin
                if [ -f "/etc/systemd/system/getty@tty1.service.d/autologin.conf" ] ; then
                    sudo rm /etc/systemd/system/getty@tty1.service.d/autologin.conf
                fi
                # Update desktop login
                update_desktop_login "False" $USER
                REQUIRE_REBOOT=true
                ;;
            *B4*)
                sudo systemctl set-default graphical.target
                sudo ln -fs /lib/systemd/system/getty@.service /etc/systemd/system/getty.target.wants/getty@tty1.service
                # Remove autologin
                if [ -f "/etc/systemd/system/getty@tty1.service.d/autologin.conf" ] ; then
                    sudo rm /etc/systemd/system/getty@tty1.service.d/autologin.conf
                fi
                # Update desktop login
                update_desktop_login "True" $USER
                REQUIRE_REBOOT=true
                ;;
            *)
                whiptail --msgbox "Programmer error, unrecognised boot option" 20 60 2
                ;;
        esac
        if $REQUIRE_REBOOT ; then
            # Require reboot
            if [ ! -f /var/run/reboot-required ] ; then
                sudo sh -c 'echo "*** System Restart Required ***" > /var/run/reboot-required'
            fi
            # End loop
            do_finish
        fi
    fi
    # Return to main for other options
    main
}

wifi()
{
    # Reference:
    # https://github.com/rbonghi/jetson_stats/issues/1#issuecomment-437585912
    # https://askubuntu.com/questions/497395/14-04-wireless-throughput-issues
    # https://unix.stackexchange.com/questions/404287/wifi-suddenly-extremely-slow
    # Also run this to fix the Mbs rate. By default its slow so even if you add a Ethernet cable it will be terribly slow.
    # Check if iwlwifi is included
    if [ -f "/etc/modprobe.d/iwlwifi.conf" ] ; then
        local TEXT="If suddenly the wifi is extremely slow a solution to fix the Mbs rate\n"
        TEXT+="A option is enable on your wifi configuration adding:\n\n"
        TEXT+="\"options iwlwifi 11n_disable=8\" in /etc/modprobe.d/iwlwifi.conf"
        # Check if options iwlwifi 11n_disable=8 is already in /etc/modprobe.d/iwlwifi.conf
        if grep -Fxq "options iwlwifi 11n_disable=8" "/etc/modprobe.d/iwlwifi.conf" ; then
            whiptail --title "$TITLE" --msgbox "$TEXT\n\nInstalled!" 14 78
        else
            # Ask if you want add this rule
            whiptail --title "$TITLE" --yesno "$TEXT\n\nDo you want update?" 14 78 2
            if [ $? -eq 0 ]; then # yes
                # Add 11n_disable=8 in iwlwifi.conf
                echo "options iwlwifi 11n_disable=8" | sudo tee -a /etc/modprobe.d/iwlwifi.conf
                # Require reboot
                if [ ! -f /var/run/reboot-required ] ; then
                    sudo sh -c 'echo "*** System Restart Required ***" > /var/run/reboot-required'
                fi
            fi
        fi
    else
        whiptail --title "$TITLE" --msgbox "No wifi module recognize\n/etc/modprobe.d/iwlwifi.conf does not exist." 8 78
    fi
    # Return to main menu
    main
}

update()
{
    # Run the installer
    sudo -H pip install --no-cache-dir -U jetson_stats
    # Show last version jtop
    whiptail --title "$TITLE" --msgbox "Jtop installed:\n - $(jtop -v 2>&1)" 8 78
    # Move to health page
    health
}

about()
{
    whiptail  --title "$TITLE" --msgbox "\
This tool provides a straight-forward way of doing initial
configuration of your NVIDIA jetson. Although it can be run
at any time, some of the options may have difficulties if
you have heavily customised your installation.\n\
\nScripts:\n - jtop\n - jetson_config\n - jetson_release\n - jetson_swap\
\nServices:\n - jetson_stats\
" 20 70 1
    # Return to main menu
    main
}


main()
{
    local MENU_LIST=("health" "Check the status of jetson-stats"\
                     "desktop" "Enable/Disable boot from desktop"\
                     "wifi" "Improve wireless throughput"\
                     "update" "Update this tool to the latest version"\
                     "about"  "Information about this configuration tool")
    # Evaluate the size
    local MENU_LIST_LEN
    let MENU_LIST_LEN=${#MENU_LIST[@]}/2
    local MENU_TEXT="$JETSON_MACHINE\n"
    MENU_TEXT+="$(jtop -v 2>&1)"
    # Make jetson_stats menu
    local OUTPUT
    OUTPUT=$(whiptail --title "$TITLE" --ok-button "Select" --cancel-button "Finish" --menu "$MENU_TEXT" 15 78 $MENU_LIST_LEN "${MENU_LIST[@]}" 3>&1 1>&2 2>&3)
    # Read exit status
    if [ $?=0 ]; then
        # Run the function with the same name
        $OUTPUT
    fi
}

do_finish() {
    # Reference:
    # https://askubuntu.com/questions/921541/how-to-determine-if-a-restart-is-needed-on-my-ubuntu-server
    if [ -f /var/run/reboot-required ] ; then
        if $INTERACTIVE ; then
            whiptail --title "$TITLE" --yes-button "Reboot" --yesno "Would you like to reboot now?" 12 60 2
            if [ $? -eq 0 ]; then # yes
                sync
                sudo reboot
            fi
        fi
        echo "${red}Require reboot!${reset}"
    fi
    exit 0
}

do_password()
{
    # https://superuser.com/questions/553932/how-to-check-if-i-have-sudo-access
    prompt=$(sudo -nv 2>&1)
    if [ $? -eq 0 ]; then
        # exit code of sudo-command is 0
        # echo "has_sudo__pass_set"
        return
    elif echo $prompt | grep -q '^sudo:'; then
        #echo "has_sudo__needs_pass"
        if $INTERACTIVE ; then
            # Open dialog menu to ask the password
            local PASSWORD=$(whiptail --passwordbox "please enter '$USER' password" 8 78 --title "$TITLE" 3>&1 1>&2 2>&3)
            exitstatus=$?
            if [ $exitstatus = 0 ]; then
                # Check if the root password if good
                if $(echo "$PASSWORD" | sudo -S -i true) ; then
                    echo 
                fi
                do_password
            else
                # Exit from request
                exit 0
            fi
        else
            # Request sudo password
            sudo -v
        fi
    else
        echo "no sudo"
        return
    fi
}

usage()
{
    if [ "$1" != "" ]; then
        echo "${red}$1${reset}"
    fi

    echo "jetson_config. [-p [OPTION]] | [--uninstall] | [-h]"
    echo "   Configurator for your NVIDIA Jetson and health check for jetson-stats."
    echo "Usage:"
    echo "$0 [options]"
    echo "options,"
    echo "   -h|--help    | This help"
    echo "   -p|--option  | Select page to open the script"
    echo "   -u|--user    | Change user to run this script"
    echo "   --install    | Run installer. Equivalent to fix jetson-stats in health"
    echo "   --uninstall  | Run uninstaller for all packages"
}

loop()
{
    local FUNCTION=""
    local INSTALLER=false
    local UNINSTALLER=false
    local PIP_INSTALLER=false
    # Decode all information from startup
    while [ -n "$1" ]; do
        case "$1" in
            -h|--help)
                usage
                exit 0
                ;;
            -p|--option)
                FUNCTION=$2
                shift 1
                ;;
            -u|--users)
                USER=$2
                shift 1
                ;;
            --install)
                FORCE=true
                INTERACTIVE=false
                INSTALLER=true
                ;;
            --uninstall)
                INTERACTIVE=false
                UNINSTALLER=true
                ;;
            *)
                usage "[ERROR] Unknown option: $1"
                exit 1
            ;;
        esac
            shift 1
    done
    # Ask sudo password before run
    do_password
    if $INSTALLER ; then
        # Run the installer
        installer "${HEALTH_FUNCTIONS[@]}"
    elif $UNINSTALLER ; then
        # Clean jtop configuration and restore
        sudo jtop --restore
        # Run the uninstaller
        uninstaller "${HEALTH_FUNCTIONS[@]}"
    else
        # If exist argument pass directly the name
        if [ ! -z "$FUNCTION" ] ; then
            if type $FUNCTION | grep -i function > /dev/null; then
                # Run function
                $FUNCTION
            fi
        else
            # Run main
            main $@
        fi
    fi
}

loop $@
exit 0
# EOF
