#!/bin/bash

# This script has three parts which all use the Substrate runtime:
# - Pallet benchmarking to update the pallet weights
# - Overhead benchmarking for the Extrinsic and Block weight
# - Machine benchmarking

# Initialize variables at the start
skip_build=false
start_pallet=""

while getopts 'bfp:v' flag; do
    case "${flag}" in
        b)
            # skip build
            skip_build=true;;
        f)
            # fail in any sub-command in a pipe fails
            set -o  pipefail
            # fail of undeclared variables
            set -u
            #fail if any sub-command fails
            set -e
            # fail on traps
            set -E
            ;;
        p)
            # start at pallet
            start_pallet="${OPTARG}"
            ;;
        v)
            # echo all executed commands
            set -x
            ;;
        *)
            # exit early
            echo "Bad options. Chech script"
            exit 1
            ;;
    esac
done

ERR_FILE="benchmarking_errors.txt"
EXECUTABLE=./target/release/ghost
RUNTIME="casper"

rm -f $ERR_FILE

if [ "$skip_build" != true ]
then
    echo "[+] Compiling Node with 'runtime-benchmarks' features"
    cargo build --release --locked --features=runtime-benchmarks,$RUNTIME-native
fi

# Update the block and extrinsic overhead weights
echo "[+] Benchmarking block and extrinsic overheads..."
OUTPUT=$(
    $EXECUTABLE benchmark overhead \
    --chain="${RUNTIME}-local" \
    --wasm-execution=compiled \
    --weight-path="./runtime/${RUNTIME}/constants/src/weights/" \
    --header=./file_header.txt \
    --warmup=10 \
    --repeat=100 2>&1
)
if [ $? -ne 0 ]; then
    echo "$OUTPUT" >> "$ERR_FILE"
    echo "[-] Failed to benchmark the block and extrinsic weight. Error written to $ERR_FILE; continuing..."
fi

echo "[+] Benchmarking the machine..."
OUTPUT=$(
    $EXECUTABLE benchmark machine --allow-fail --chain=${RUNTIME}-dev 2>&1
)
if [ $? -ne 0 ]; then
    # Do not write the error file since it is not a benchmarking error
    echo "[-] Failed the machine benchmark:\n$OUTPUT"
fi

echo "[+] Benchmarking the RocksDb storage..."
OUTPUT=$(
    $EXECUTABLE benchmark storage \
    --chain="${RUNTIME}-dev" \
    --state-version=0 \
    --mul=1.1 \
    --weight-path="./runtime/${RUNTIME}/constants/src/weights/" \
    --header=./file_header.txt 2>&1
)
if [ $? -ne 0 ]; then
    echo "$OUTPUT" >> "$ERR_FILE"
    echo "[-] Failed to benchmark the RocksDb storage. Error written to $ERR_FILE; continuing..."
fi

echo "[+] Benchmarking the ParityDb storage..."
OUTPUT=$(
    $EXECUTABLE benchmark storage \
    --chain="${RUNTIME}-dev" \
    --state-version=0 \
    --mul=1.1 \
    --database=paritydb \
    --weight-path="./runtime/${RUNTIME}/constants/src/weights/" \
    --header=./file_header.txt 2>&1
)
if [ $? -ne 0 ]; then
    echo "$OUTPUT" >> "$ERR_FILE"
    echo "[-] Failed to benchmark the ParityDb storage. Error written to $ERR_FILE; continuing..."
fi

EXCLUDED_PALLETS=(
    # helper pallets
    "frame_benchmarking"
)

# Load all pallet names in an array
ALL_PALLETS=($(
    $EXECUTABLE benchmark pallet --list=pallets --no-csv-header --chain=${RUNTIME}-dev 2>/dev/null | grep -v ':' | grep -v 'baseline'
))

# Filter out the excluded pallets by concatenating the arrays and discarding duplicates
PALLETS=($(
    printf '%s\n' "${ALL_PALLETS[@]}" | grep -Fvxf <(printf '%s\n' "${EXCLUDED_PALLETS[@]}")
))
TOTAL_PALLETS=${#PALLETS[@]}
CURRENT=0

echo "[+] Benchmarking ${#PALLETS[@]} pallets by excluding ${#EXCLUDED_PALLETS[@]} from ${#ALL_PALLETS[@]}"

for PALLET in "${PALLETS[@]}"; do
    CURRENT=$((CURRENT + 1))

    # If `-p` is used, skip benchmark until the start pallet.
    if [ ! -z "$start_pallet" ] && [ "$start_pallet" != "$PALLET" ]
    then
        echo "[+] ($CURRENT/$TOTAL_PALLETS) Skipping ${PALLET}..."
        continue
    else
        unset start_pallet
    fi

    if [[ -z "$PALLET" || "$PALLET" == "::" || "$PALLET" == "baseline" || "$PALLET" == "pallet" ]]; then
        echo "[+] ($CURRENT/$TOTAL_PALLETS) Skipping garbage names: '${PALLET}'"
        continue
    fi

    FOLDER="$(echo "${PALLET#*_}" | tr '_' '-')";
    WEIGHT_FILE="./runtime/${RUNTIME}/src/weights/${PALLET}.rs"
    echo "[+] Benchmarking $PALLET for $RUNTIME with weight file $WEIGHT_FILE ($CURRENT/$TOTAL_PALLETS)";

    OUTPUT=$(
        $EXECUTABLE benchmark pallet \
        --chain="${RUNTIME}-dev" \
        --steps=50 \
        --repeat=20 \
        --pallet="$PALLET" \
        --extrinsic="*" \
        --wasm-execution=compiled \
        --heap-pages=4096 \
        --header=./file_header.txt \
        --output=${WEIGHT_FILE} 2>&1
    )

    if [ $? -ne 0 ]; then
        echo "$OUTPUT" >> "$ERR_FILE"
        echo "[-] Failed to benchmark $PALLET. Error written to $ERR_FILE; continuing..."
    fi
done

# Check if the error file exists
if [ -f "$ERR_FILE" ]; then
    echo "[-] Some benchmarks failed. See: $ERR_FILE"
    exit 1
else
    echo "[+] All benchmarks passed"
    exit 0
fi
