#!/bin/bash

# Change the current directory to the script directory
BASEDIR=$(dirname "$0")
cd $BASEDIR
# Look for the local .env file
while [[ "$PWD" != "$HOME" && ! -e .env ]]; do cd ..; done
if [[ -e .env ]]; then
  set -a; source .env; set +a
  echo "Using .env file found in $PWD"
else
  echo "CANCELED"
  echo "  No .env file found in the directory tree."
  echo "  Please setup the .env file and retry"
  exit 1
fi

# Change the current directory to the home directory
cd $INFAMOUSDIR
echo "EXECUTED FROM $PWD"

if [[ $# -lt 1 ]]; then
    echo " "
    echo "CANCELED"
    echo "  No arguments provided."
    echo "  Syntax: ./tools/run_on_apiserver.sh remote_path [log_subdir] [env_name] [-- script_args...]"
    echo " "
    echo "  Examples:"
    echo " "
    echo "      To run script.py located in INFAMOUSDIR/path/to on the api server:"
    echo "      ./tools/run_on_apiserver.sh path/to/script.py"
    echo " "
    echo "          Only pass the path from the remote INFAMOUSDIR to the file!"
    echo "          The path to the remote INFAMOUSDIR will be read from the server-side .env file."
    echo " "
    echo "          By default, python files are executed in an environment called oneoff-env,"
    echo "          and the log is piped into a subdirectory called oneoff"
    echo " "
    echo "      If you want to pipe the log into a specific subdirectory, you can pass"
    echo "      the name of the directory as a secondary argument:"
    echo "      ./tools/run_on_apiserver.sh path/to/script.py myfancylog"
    echo " "
    echo "      If you want to run script.py in a specific environment, you can pass"
    echo "      the name of the environment as a tertiary argument:"
    echo "      ./tools/run_on_apiserver.sh path/to/script.py myfancylog myfancy-env"
    echo " "
    echo "      To pass arguments to the script, use '--' followed by the script arguments:"
    echo "      ./tools/run_on_apiserver.sh path/to/script.py myfancylog myfancy-env -- arg1 arg2"
    echo " "
    echo "      You can also run shell files that trigger multiple python files, notebooks etc."
    echo "      In such case, you will need to activate the desired env in your shell script."
    echo "      ./tools/run_on_apiserver.sh path/to/script.sh"
    echo " "
    echo "      Specifying the log subdirectory is also possible for shell files"
    echo "      ./tools/run_on_apiserver.sh path/to/script.sh myfancylog"
    exit 1
fi

# Parse command line arguments
SCRIPT_PATH=$1
shift
LOGSUBDIR="oneoff"
ENV_NAME="oneoff-env"
SCRIPT_ARGS=""

# Handle optional arguments and '--'
while [[ $# -gt 0 ]]; do
    case $1 in
        --)
            shift
            SCRIPT_ARGS="$@"
            break
            ;;
        *)
            if [[ -z $LOGSUBDIR_SET ]]; then
                LOGSUBDIR=$1
                LOGSUBDIR_SET=true
            elif [[ -z $ENV_NAME_SET ]]; then
                ENV_NAME=$1
                ENV_NAME_SET=true
            fi
            shift
            ;;
    esac
done

ssh -i "$pem_path" "$api_user@$bastion_server" ssh "$api_user@$api_server" << EOF

    sudo -u "$api_suuser" bash -c '

        cd \$HOME

        # Look for the remote .env file
        while [[ "\$PWD" != "\$HOME" && ! -e .env ]]; do cd ..; done
        if [[ -e .env ]]; then
            set -a; source .env; set +a
            echo "Using remote .env file found in \$PWD"
        else
            echo "CANCELED"
            echo "  No .env file found in the remote directory tree."
            echo "  Please setup the .env file and retry"
            exit 1
        fi

        # Define remote paths
        FILEPATH=\$INFAMOUSDIR/$SCRIPT_PATH
        LOGPATH=\$LOGDIR/$LOGSUBDIR
        ENVPATH=\$ENVDIR/$ENV_NAME
        FILETYPE=${SCRIPT_PATH##*.}

        if [ ! -e \$FILEPATH ]; then
            echo " "
            echo "ERROR"
            echo "  \$FILEPATH does not exist"
            echo "  Log on and check if the file is where it should be, and upload the file if needed"
            exit 1
        fi
        if [ ! -e \$LOGPATH ]; then
            echo " "
            echo "ERROR"
            echo "  \$LOGPATH does not exist!"
            echo "  Please create a log directory for oneoff scripts or pass an existing subdirectory"
            exit 1
        fi
        if [ ! -e \$ENVPATH ] && [ \$FILETYPE = "py" ]; then
            echo " "
            echo "ERROR"
            echo "  \$ENVPATH does not exist!"
            echo "  Please create an environment for oneoff scripts, pass an existing environment,"
            echo "  or pass a path to a shell script"
            exit 1
        fi

        # Determine script type based on extension
        case \$FILETYPE in
            py)
                source \$ENVPATH/bin/activate || {
                    echo " "
                    echo "ERROR"
                    echo "  Failed to activate \$ENVPATH"
                    exit 1
                }
                
                nohup python3.10 "\$FILEPATH" $SCRIPT_ARGS > "\$LOGPATH/\$(basename "$SCRIPT_PATH").log" 2>&1 &
                ;;
            sh)
                chmod +x \$FILEPATH
                nohup "\$FILEPATH" $SCRIPT_ARGS > "\$LOGPATH/\$(basename "$SCRIPT_PATH").log" 2>&1 &
                ;;
            *)
                echo " "
                echo "ERROR"
                echo "  Unsupported file type: \$FILETYPE"
                echo "  Only .py and .sh files are supported."
                exit 1
                ;;
        esac

        exit 0
    '
EOF

if [ $? -ne 0 ]; then
    echo "CANCELED"
    echo "  Check the errors above."
    exit 1
fi

ssh -i $pem_path $api_user@$bastion_server ssh $api_user@$api_server << EOF

    sudo -u "$api_suuser" bash -c '

        cd \$HOME

        # Look for the remote .env file
        while [[ "\$PWD" != "\$HOME" && ! -e .env ]]; do cd ..; done
        if [[ -e .env ]]; then
            set -a; source .env; set +a
            echo "Using remote .env file found in \$PWD"
        else
            echo "CANCELED"
            echo "  No .env file found in the remote directory tree."
            echo "  Please setup the .env file and retry"
            exit 1
        fi

        FILEPATH=\$INFAMOUSDIR/$SCRIPT_PATH

        if pgrep -f \$FILEPATH >/dev/null; then
            echo "\$FILEPATH is running."
        else
            echo "Process has terminated."
        fi

        exit
    '
EOF