#!/usr/bin/env bash
# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# This software product is a proprietary product of Nvidia Corporation and
# its affiliates (the "Company") and all right, title, and interest in and
# to the software product, including all associated intellectual property
# rights, are and shall remain exclusively with the Company.
#
# This software product is governed by the End User License Agreement
# provided with the software product.

set -uo pipefail

# allow functions to return 1 without indicating an error
set +e

UDS_PATH="$HOME/.local/share/nvue/nvd-socket"
declare -r FSM_LIMBO=0 FSM_STDOUT=1 FSM_STDERR=2 FSM_EXIT_CODE=3


is_tab_complete() {
    # return 1 if we are servicing a TAB completion
    local cmd=("$@")

    # if the user did "nv<ENTER>" the cmd array will be empty
    if [ ${#cmd[@]} -eq 0 ]; then
        return 0
    fi

    tab_options=("-t" "--tab")

    for arg in "${cmd[@]}"; do
        for option in "${tab_options[@]}"; do
            if [[ "$arg" == "$option" ]]; then
                return 1
            fi
        done
    done

    return 0
}

use_nv_via_nvd() {
    # return 1 if this is a command that is supported by nvd
    local cmd=("$@")

    # if the user did "nv<ENTER>" the cmd array will be empty
    if [ ${#cmd[@]} -eq 0 ]; then
        return 0
    fi

    # "nv set"
    if [[ ${cmd[0]} == se* || ${cmd[0]} == s ]]; then
        return 1
    fi

    # nv unset
    if [[ ${cmd[0]} == u* ]]; then
        return 1
    fi

    # TAB complete, help, etc
    supported_options=("-h" "--help" "-t" "--tab" "-q" "--brief-help")

    for arg in "${cmd[@]}"; do
        for option in "${supported_options[@]}"; do
            if [[ "$arg" == "$option" ]]; then
                return 1
            fi
        done
    done

    return 0
}

nvd_running() {
    # return 1 if there is a nvd.py process that we own
    ps -fjH -u $USER | grep nvd.py | grep -v grep > /dev/null

    if [[ $? -eq 0 ]]; then
        return 1
    else
        return 0
    fi
}

start_nvd() {
    nvd_running
    nvd_running_result=$?

    if [[ $nvd_running_result -eq 0 ]]; then
        # nvd is not running, start it
        nvd.py --quiet-period "${NVUE_NVD_QUIET_PERIOD:-300}" --daemon 2>&1 > /dev/null
    else
        if [ ! -e "$UDS_PATH" ]; then
            # nvd is running but UDS_PATH is not there, kill nvd. This would be a pretty
            # contrived scenario where someone manually deleted the UDS_PATH.
            nvd_pid=`ps -fjH -u $USER | grep nvd.py | grep -v grep | awk -F ' ' '{print $2}'`

            if [ "$nvd_pid" != "" ]; then
                # kill nvd and then sleep for 1s to verify it is dead.  If it isn't dead, do a kill -9
                kill $nvd_pid
                sleep 1
                nvd_pid=`ps -fjH -u $USER | grep nvd.py | grep -v grep | awk -F ' ' '{print $2}'`

                if [ "$nvd_pid" != "" ]; then
                    kill -9 $nvd_pid
                    sleep 1
                fi
            fi

            nvd.py --quiet-period "${NVUE_NVD_QUIET_PERIOD:-300}" --daemon 2>&1 > /dev/null
        fi
    fi
}

# send data to the nvd socket
send_to_socket() {
    output=$(printf "%s\n" "$1" | nc -U "$UDS_PATH" 2>/dev/null)
}

parse_nvd_response() {
    # Process the response
    state=$FSM_LIMBO
    exit_code=1

    while IFS= read -r line || [[ -n $line ]]; do
        case "$state" in
            $FSM_STDOUT)
                case "$line" in
                    "</STDOUT>") state=$FSM_LIMBO ;;
                    *) echo "$line" ;;
                esac
                ;;
            $FSM_STDERR)
                case "$line" in
                    "</STDERR>") state=$FSM_LIMBO ;;
                    *) echo "$line" >&2 ;;
                esac
                ;;
            $FSM_EXIT_CODE)
                case "$line" in
                    "</EXIT_CODE>") state=$FSM_LIMBO ;;
                    *) exit_code=$line ;;
                esac
                ;;
            $FSM_LIMBO)
                case "$line" in
                    "<STDOUT>") state=$FSM_STDOUT ;;
                    "<STDERR>") state=$FSM_STDERR ;;
                    "<EXIT_CODE>") state=$FSM_EXIT_CODE ;;
                    *) echo "ERROR: invalid line: $line" >&2 ;;
                esac
                ;;
        esac
    done < <(printf '%s' "$output")
    return "${exit_code}"
}


args_to_json() {
    local arr=("$@")
    local json="[\"nv\","
    local first=1

    for item in "${arr[@]}"; do
        if [ "$first" -eq 1 ]; then
            first=0
        else
            json+=","
        fi
        json+="\"${item//\"/\\\"}\""
    done

    json+="]"
    echo "$json"
}


main() {

    # Some commands must be interactive and cannot use the faster nvd approach
    # Check to see if this command is elligible for nvd
    use_nv_via_nvd "$@"
    use_nv_via_nvd_result=$?

    # use nvd
    if [[ $use_nv_via_nvd_result -eq 1 ]]; then

        if [ -e "$UDS_PATH" ]; then
            # create a JSON string of our command line args
            JSON=$(args_to_json "$@")

            # send our command line args to nvd via a unix domain socket
            send_to_socket "$JSON"
            send_to_socket_status=$?
        else
            send_to_socket_status=1
        fi

        # If we were able to connect to nvd then great, we will exit with the exit_code
        # we received from nvd.  If not though we will fall back to nv-interactive.
        if [[ $send_to_socket_status -eq 0 ]]; then
            parse_nvd_response "$@"
            exit_code=$?

            # 199 is a special code used by nvd to signal to the nv client that
            # it must use nv-interactive.  This is a corner case that only
            # happens if the "nv set" command must drop into a prompt to get
            # some data from the user.  This is typically used for prompting
            # the user for a password.
            if [[ $exit_code != "199" ]]; then
                exit $exit_code
            fi
        else
            # We failed to connect to nvd.  This is normal if nvd.py isn't running and
            # is ok, we will fall back to nv-interactive in this scenario.
            #
            # If nv ran for something other than TAB complete then start nvd in the background.
            # We do not start nvd for TAB complete because bash locks up in that scenario.
            is_tab_complete $@
            is_tab_complete_result=$?

            if [[ $is_tab_complete_result -eq 0 ]]; then
                start_nvd
            fi
        fi
    fi

    nv-interactive "$@"
    exit_code=$?
    exit $exit_code
}

main "$@"
