#! /bin/bash

# This script has two different modes of operation:
# if KOKKOS_NEXTSILICON_TEST_TELEMETRYLESS is anything, we run in a telemetry-less mode, where any projection error counts as a failure
# else, we run the two-pass telemetry-based flow

set -e

# Skip nextsystemd entirely if no tests would run (respects --gtest_filter and similar args).
# For GTest binaries, count test cases listed by --gtest_list_tests (indented two spaces).
if ./"$1" --help 2>&1 | grep -q 'This program contains tests written using Google Test'; then
    _test_count=$(./"$1" "${@:2}" --gtest_list_tests 2>/dev/null | grep -c '^  ' || true)
    if [ "${_test_count}" -eq 0 ]; then
        echo "nextsilicon-test-wrapper.sh: no tests selected, skipping"
        exit 0
    fi
fi

# make sure nextsystemd shuts down no matter what
cleanup() {
        # code from last thing that ran
        local result=$?
        echo "cleanup..."
        nextcli system terminate || true
        wait ${NEXTSYSTEMD_PID}
        echo "exit with $result"
        exit $result
}
trap cleanup EXIT

patch_dir=$(mktemp -d)

# Kill any lingering NS processes left by a previously timed-out test
pkill -x nextsystemd   2>/dev/null || true
pkill -x nextoptd      2>/dev/null || true
pkill -x nextruntimed  2>/dev/null || true
pkill -f "webapps/collector/influxdb" 2>/dev/null || true
rm -f "/run/user/$(id -u)/nextsystemd/observability"/*.sock 2>/dev/null || true
sleep 2

if [[ -n "${KOKKOS_NEXTSILICON_TEST_TELEMETRYLESS:-}" ]]; then
    echo "nextsilicon-test-wrapper.sh: KOKKOS_NEXTSILICON_TEST_TELEMETRYLESS is set"

    # configure nextsystemd
    tee ${patch_dir}/kokkos.patch <<EOF
optimizer-pi:
  ebbm-enabled: true
  max-automatic-migration-move-count: 18446744073709551615 # MAX_UINT64
  # switch to telemetry-less mode
  enable-telemetry-less: true
projection:
  # increase cache size to avoid eviction of projection data
  cache-max-size-mib: 8192
  # Only accept single CG projection
  cg-split: single-cg
telem:
  # disable device telemetry
  mode: none
  enabled-telemetries: ""
EOF

    # start nextsystemd without starting UI tools (device telemetry is disabled
    # via the telem: block above)
    nextsystemd --ui-collector-address none --cfg-file ${patch_dir}/kokkos.patch &
    NEXTSYSTEMD_PID=$!

    # the following nextcli command seems to time out too aggressively on 1.3.0-66
    sleep 20

    # verify that nextsystemd is running
    nextcli system status | grep -q 'Service: UP' || exit 127

    # telemetry-less run
    nextloader -- ./"$1" "${@:2}"
else
    echo "nextsilicon-test-wrapper.sh: KOKKOS_NEXTSILICON_TEST_TELEMETRYLESS is not set"

    # configure nextsystemd
    tee ${patch_dir}/kokkos.patch <<EOF
optimizer-pi:
  ebbm-enabled: true
  max-automatic-migration-move-count: 18446744073709551615 # MAX_UINT64
  # FIXME_NEXTSILICON: use one tile, some unit test binaries cause configuration error for full device
  numa-regions-per-device: 1
  numa-regions-width: 1
  numa-regions-height: 1
  mlc:
    # try to offload every parallel region
    acceleration-threshold: 1
projection:
  # increase cache size to avoid eviction of projection data
  cache-max-size-mib: 8192
  # Only accept single CG projection
  cg-split: single-cg
telem:
  # disable device telemetry
  mode: none
  enabled-telemetries: ""
EOF

    # start nextsystemd without starting UI tools (device telemetry is disabled
    # via the telem: block above)
    nextsystemd --ui-collector-address none --cfg-file ${patch_dir}/kokkos.patch &
    NEXTSYSTEMD_PID=$!

    # the following nextcli command seems to time out too aggressively on 1.3.0-66
    sleep 20

    # verify that nextsystemd is running
    nextcli system status | grep -q 'Service: UP' || exit 127

    # training run
    nextloader -- ./"$1" "${@:2}"

    # wait for optimization to start
    sleep 15

    # wait for optimization/projection to finish, up to 1 hour
    SECONDS=0
    while [ $SECONDS -lt 3600 ]; do
        # check current state
        ret=0
        status="$(nextcli application status | grep 'Optimization state:')"
        if [[ $status == *IDLE* ]]; then
            # no mills found, return
            exit 0
        elif [[ $status == *IMPROVED* ]]; then
            # optimization/projection finished, do device run
            nextloader -- ./"$1" "${@:2}"
            exit
        elif [[ $status == *OPTIMIZING* ]]; then
            # optimization/projection still running, wait 30 more seconds
            sleep 30
            continue
        else
            # in some other state, something went wrong
            exit 2
        fi
    done

    echo "error: timed out" >&2
    exit 127
fi