#!/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

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

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

EXECUTABLE=./target/release/ghost
RUNTIME="casper"

EXCLUDED_PALLETS=(
    # helper pallets
    "pallet_election_provider_support_benchmarking"
    # pallets without automatic benchmarking
    # "pallet_babe"
    # "pallet_grandpa"
    "pallet_bags_list"
    "frame_system"
    "pallet_staking"
    "pallet_offences"
)

# Load all pallet names in an array
ALL_PALLETS=($(
    $EXECUTABLE benchmark pallet --list=pallets --no-csv-header --chain=${RUNTIME}-dev
))

# Filter out the excluded pallets by concatenating the arrays and discarding duplicates
PALLETS=($({ printf '%s\n' "${ALL_PALLETS[@]}" "${EXCLUDED_PALLETS[@]}"; } | sort | uniq -u))

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

# Define the error file
ERR_FILE="benchmarking_errors.txt"
# Delete the error file before each one
rm -f $ERR_FILE

for PALLET in "${PALLETS[@]}"; do
    # If `-p` is used, skip benchmark until the start pallet.
    if [ ! -z "$start_pallet" ] && [ "$start_pallet" != "$PALLET" ]
    then
        echo "[+] Skipping ${PALLET}..."
        continue
    else
        unset start_pallet
    fi

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

    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

# # Update the block and extrinsic overhead weights
# echo "[+] Benchmarking block and extrinsic overheads..."
# OUTPUT=$(
#     $EXECUTABLE benchmark overhead \
#     --chain="${RUNTIME}-dev" \
#     --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 --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

# 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