#!/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`

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"\
                  "jetson_performance" "Status jetson clock service"\
                  "jetson_stats_up" "Jetson-stats startup 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
}

status_jetson_performance()
{
    # Check if jetson_performance is executable
    if [ ! -x "$JETSON_STATS_FOLDER/jetson_performance.sh" ] ; then
        echo 1
        return
    fi
    # Link jetson_performance in list
    if [ ! -f "/etc/systemd/system/jetson_performance.service" ] ; then
        echo 1
        return
    fi
    # Check if file are same
    if ! cmp --silent "/etc/systemd/system/jetson_performance.service" "$JETSON_STATS_FOLDER/jetson_performance.service" ; then
        echo 1
        return
    fi
    # Check if jetson performance is correctly installed
    local JE_PERFOMANCE_STATUS="$(systemctl is-active jetson_performance.service)"
    if [ $JE_PERFOMANCE_STATUS = "failed" ] ; then
        echo 1
        return
    fi
    # Everything is fine
    echo 0
}

fix_jetson_performance()
{
    # Make executable
    sudo chmod +x "$JETSON_STATS_FOLDER/jetson_performance.sh"
    # Check link jetson_performance in list
    if [ ! -f "/etc/systemd/system/jetson_performance.service" ] ; then
        # Using a copy of the service to prevent the remotion when is call the function "disable"
        sudo cp "$JETSON_STATS_FOLDER/jetson_performance.service" "/etc/systemd/system/jetson_performance.service"
        # sudo ln -s "$JETSON_STATS_FOLDER/jetson_performance.service" "/etc/systemd/system/jetson_performance.service"
    fi
    # Update service list
    sudo systemctl daemon-reload
}

uninstall_jetson_performance()
{
    # 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
    if [ -f "/etc/systemd/system/jetson_performance.service" ] ; then
        sudo rm "/etc/systemd/system/jetson_performance.service"
    fi
    # Update service list
    sudo systemctl daemon-reload
}

jetson_performance()
{
    local TEXT="Status jetson_performance. The jetson_clock service.\n"
    # Write status jetson_performance service
    if [ -x "$JETSON_STATS_FOLDER/jetson_performance.sh" ] ; then
        TEXT+=" - jetson_performance.sh is executable\n"
    else
        TEXT+=" - jetson_performance.sh is NOT executable\n"
    fi
    # Status service
    local JE_PERFOMANCE_STATUS="$(systemctl is-active jetson_performance.service)"
    TEXT+=" - [$JE_PERFOMANCE_STATUS] jetson_performance service\n"
    # Enabled service
    local JE_PERFOMANCE_ENABLED=$(systemctl is-enabled jetson_performance.service 2> /dev/null)
    if [ -z $JE_PERFOMANCE_ENABLED ] ; then
        JE_PERFOMANCE_ENABLED="failed"
    fi
    TEXT+=" - [$JE_PERFOMANCE_ENABLED] jetson_performance enable"
    # Check status jetson_performance
    if [ $(status_jetson_performance) -eq 1 ] ; then
        # Fix jetson_performance
        if (whiptail --title "$TITLE" --yesno "$TEXT\n\nDo you want fix?" 11 78); then
            fix_jetson_performance
        fi
    else
        # Run Message
        whiptail --title "$TITLE" --msgbox "$TEXT" 10 78
    fi
    # Return to healt function
    health
}

status_jetson_stats_up()
{
    # Remove old jetson_fan service
    if [ -f "/etc/systemd/system/jetson_fan.service" ] ; then
        echo 1
        return
    fi
    # Check status file
    if [ ! -x "$JETSON_STATS_FOLDER/jetson_stats_boot.sh" ] ; then
        echo 1
        return
    fi
    # Check link jetson_stats_boot in list
    if [ ! -f "/etc/systemd/system/jetson_stats_boot.service" ] ; then
        echo 1
        return
    fi
    # Check if file are same
    if ! cmp --silent "/etc/systemd/system/jetson_stats_boot.service" "$JETSON_STATS_FOLDER/jetson_stats_boot.service" ; then
        echo 1
        return
    fi
    # Check status service
    local JE_FAN_STATUS="$(systemctl is-active jetson_stats_boot.service)"
    if [ $JE_FAN_STATUS = "failed" ] ; then
        echo 1
        return
    fi
    local JE_FAN_ENABLED=$(systemctl is-enabled jetson_stats_boot.service 2> /dev/null)
    if [ -z $JE_FAN_ENABLED ] ; then
        JE_FAN_ENABLED="failed"
        echo 1
        return
    fi
    # Everything is fine
    echo 0
}

fix_jetson_stats_up()
{
    # Remove old jetson_fan service
    if [ -f "/etc/systemd/system/jetson_fan.service" ] ; then
        # Disable service
        sudo systemctl disable jetson_fan.service
        # Remove fan service from /etc/init.d
        sudo rm "/etc/systemd/system/jetson_fan.service"
        # Update service list
        sudo systemctl daemon-reload
    fi
    # Make executable
    sudo chmod +x "$JETSON_STATS_FOLDER/jetson_stats_boot.sh"
    # Link jetson start up service
    if [ ! -f "/etc/systemd/system/jetson_stats_boot.service" ] ; then
        # Using a copy of the service to prevent the remotion when is call the function "disable"
        sudo cp "$JETSON_STATS_FOLDER/jetson_stats_boot.service" "/etc/systemd/system/jetson_stats_boot.service"
        # Old configuration
        # sudo ln -s "$JETSON_STATS_FOLDER/jetson_stats_boot.service" "/etc/systemd/system/jetson_stats_boot.service" 2> /dev/null
    fi
    # Update service list
    sudo systemctl daemon-reload
    # Enable jetson fan at startup
    sudo systemctl enable jetson_stats_boot.service
}

uninstall_jetson_stats_up()
{
    # Disable service
    sudo systemctl disable jetson_stats_boot.service
    # Remove fan service from /etc/init.d
    if [ -f "/etc/systemd/system/jetson_stats_boot.service" ] ; then
        sudo rm "/etc/systemd/system/jetson_stats_boot.service"
    fi
    # Update service list
    sudo systemctl daemon-reload
}

jetson_stats_up()
{
    local TEXT="Status jetson_stats_boot. The jetson_stats start-up service.\n"
    if [ -x "$JETSON_STATS_FOLDER/jetson_stats_boot.sh" ] ; then
        TEXT+=" - jetson_stats_boot.sh is executable\n"
    else
        TEXT+=" - jetson_stats_boot.sh is NOT executable\n"
    fi
    # Status service
    local JE_START_UP_STATUS="$(systemctl is-active jetson_stats_boot.service)"
    TEXT+=" - [$JE_START_UP_STATUS] jetson_stats_boot service\n"
    # Enabled service
    local JE_FAN_ENABLED=$(systemctl is-enabled jetson_stats_boot.service 2> /dev/null)
    if [ -z $JE_FAN_ENABLED ] ; then
        JE_FAN_ENABLED="failed"
    fi
    TEXT+=" - [$JE_FAN_ENABLED] jetson_stats_boot enable"
    # Check status jetson_performance
    if [ $(status_jetson_stats_up) -eq 1 ] ; then
        # Fix jetson_performance
        if (whiptail --title "$TITLE" --yesno "$TEXT\n\nDo you want fix?" 11 78); then
            fix_jetson_stats_up
        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 jetson_performance
    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
}

jetson_clocks()
{
    # Reference
    # https://askubuntu.com/questions/19320/how-to-enable-or-disable-services
    local MENU="jetson_clocks is control the performance of your board.\n"
    # Read jtop status
    local JC_STATUS=""
    if [ ! -z $JTOP_PYTHON_NAME ] ; then
        local JC_STATUS=$(sudo $JTOP_PYTHON_NAME -c "from jtop import jtop; print(jtop().jc.status)")
        MENU+=" - [$JC_STATUS] jetson_clocks\n"
    fi
    # Jetson_performance status
    local JE_PERFOMANCE_STATUS="$(systemctl is-active jetson_performance.service)"
    MENU+=" - [$JE_PERFOMANCE_STATUS] jetson_performance service\n"
    # Enabled service
    local JE_PERFOMANCE_ENABLED=$(systemctl is-enabled jetson_performance.service 2> /dev/null)
    if [ -z $JE_PERFOMANCE_ENABLED ] ; then
        JE_PERFOMANCE_ENABLED="failed"
    fi
    MENU+=" - [$JE_PERFOMANCE_ENABLED] jetson_performance enable"
    # Make GUI
    if $INTERACTIVE ; then
        if [ "$JC_STATUS" = "True" ] && [ "$JE_PERFOMANCE_STATUS" != "active" ] ; then
            # Run message script
            whiptail --title "$TITLE" --msgbox "$MENU\n\njetson_clocks running without jetson_performance.service!" 14 78 2
            # Return to main
            loop
        fi
        # Make jetson_clocks service options
        BOOTOPT=$(whiptail --title "$TITLE" --menu "$MENU\nOptions:" 18 78 6 \
            "<-- Back" "Return to the main menu" \
            "Start" "Activate jetson_clocks service" \
            "Stop" "Deactivate jetson_clocks service" \
            "Enable" "Enable jetson_clocks service at system boot (delay 60s)" \
            "Disable" "Disable jetson_clocks service at system boot" \
        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
                main
                ;;
            "Start") 
                # Activate jetson_clocks service
                sudo systemctl start jetson_performance.service
                # Return to jetson_clocks menu
                jetson_clocks
                ;;
            "Stop") 
                # Activate jetson_clocks service
                sudo systemctl stop jetson_performance.service
                # Return to jetson_clocks menu
                jetson_clocks
                ;;
            "Enable") 
                # Activate jetson_clocks service
                sudo systemctl enable jetson_performance.service
                # Return to jetson_clocks menu
                jetson_clocks
                ;;
            "Disable") 
                # Activate jetson_clocks service
                sudo systemctl disable jetson_performance.service
                # Return to jetson_clocks menu
                jetson_clocks
                ;;
            *)
                whiptail --msgbox "Programmer error, unrecognised boot option" 20 60 2
                return 1
                ;;
        esac
    else
        # Return to main for other options
        main
    fi
}

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_performance\n - jetson_stats_boot\
" 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"\
                     "jetson_clocks" "Enable/Disable jetson_clocks service"\
                     "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 the 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
