#!/bin/bash

ROOT_DIR="$(dirname $0)/.."

set -e
#set -x

usage()
{
    echo -ne \
        "\nUsage:" \
        "\n    $(basename $0) -?|-h|--help" \
        "\n    $(basename $0) [options] <peer>" \
        "\n" \
        "\nOptions:" \
        "\n    -h, -?, --help                  print this help message" \
        "\n    -e, --env         <name|path>   path or dir name under <root>/env" \
        "\n    -w, --work-dir    <path>        path to place output files/logs" \
        "\n    -t, --threads     <num>         number of wrk threads (default $(nproc))" \
        "\n    -c, --connections <num>         number of wrk connections (default $(nproc))" \
        "\n    -p, --proto       <http|https>  protocol to use for connections (default http)" \
        "\n    -b, --payload     <name>        payload to use in request (default 0B)" \
        "\n        --cps                       enable CPS mode instead of RPS (effectively adds option -H 'Connection: close')" \
        "\n        --duration    <num>         duration in seconds (default 30)" \
        "\n\n"
}

platform()
{
    case "$(uname -m)" in
        x86_64)  echo x86;;
        aarch64) echo arm;;
    esac
}

main()
{
    # defaults
    local env_name="default-$(platform)"
    local work_dir="$ROOT_DIR/run/bench-client"
    local threads=$(nproc)
    local connections=$(nproc)
    local proto="http"
    local payload="0B"
    local duration="30"
    local cps=0
    #local delay=1

    # parse options
    while true; do
        case "$1" in
            -h|-\?|--help)   usage; exit 0;;
            -e|--env)        env_name=$2; shift 2;;
            -w|--work-dir)   work_dir="$2"; shift 2;;
            -t|--threads)    threads=$2; shift 2;;
            -c|--connection) connections=$2; shift 2;;
            -p|--proto)      proto=$2; shift 2;;
            -b|--payload)    payload=$2; shift 2;;
            --cps)           cps=1; shift;;
            --duration)      duration=$2; shift 2;;
            #-d|--delay)      delay=$2; shift 2;;
            --) break;;
            -?*) { echo "unknown option: $1"; usage; exit 1; } >&2;;
            *) break;;
        esac
    done

    # parse args
    local -r peer=$1
    [ -z "$peer" ] && { echo "peer address is missing"; usage; exit 1; } >&2

    # print all args
    echo "args:"
    echo "    env_name    = $env_name"
    echo "    work_dir    = $work_dir"
    echo "    proto       = $proto"
    echo "    threads     = $threads"
    echo "    connections = $connections"
    echo "    payload     = $payload"
    echo "    duration    = $duration"
    echo "    cps         = $cps"
    #echo "    delay       = $delay"
    echo "    peer        = $peer"
    echo ""

    # build full path to environment dir
    # check if environment name is a full path or just a name under <root>/env
    if [ "$(basename $env_name)" = "$env_name" ]; then
        local -r env_dir="$ROOT_DIR/env/$env_name"
    else
        local -r env_dir="$env_name"
    fi
    echo "env dir is $env_dir"

    echo "setup wrk work dir..."
    mkdir -p "$work_dir"
    
    local -r wrk_cmd="$ROOT_DIR/lib/wrk-json.sh"

    # start wrk
    if (($cps)); then
        set -x
        stdbuf -oL $wrk_cmd -t "$threads" -c "$connections" -d "$duration" -H 'Connection: close' \
            "$proto://$peer/$payload.bin" |& ( trap '' INT; tee "$work_dir/wrk.out.json" )
    else
        set -x
        stdbuf -oL $wrk_cmd -t "$threads" -c "$connections" -d "$duration" \
            "$proto://$peer/$payload.bin" |& ( trap '' INT; tee "$work_dir/wrk.out.json" )
    fi
}

main "$@"
