#!/usr/bin/env bash

set -Eeuxo pipefail

usage() {
    cat <<EOF
Usage: $0 --network <name> --fork-config <fork-config.json> --working-dir <dir> [options]

Required:
  --network                   Name of the network
  --fork-config               Path to fork_config.json
  --working-dir               Working directory
  --checks                    Checks to perform (comma-separated list). Available checks:
                                config,tarballs,ledgers. Default: all

Optional:
  --precomputed-block-prefix  Precomputed block prefix
  --reference-data-dir        Reference data directory. Defaults to empty, which skips comparison with reference data.
  --cached-hardfork-data      Path to cached hardfork data directory

This script is used to validate that an installed package is correct
according to an exported fork_config.json file.

Inputs:
- The exported mainnet full config.json fork config, with all accounts
- A working directory where ledgers/configs will be created
- Installed MINA_EXE (default: mina) and MINA_GENESIS_EXE (default: mina-create-genesis) programs
- Installed MINA_LEGACY_GENESIS_EXE program (default: mina-legacy-create-genesis)
- PACKAGED_DAEMON_CONFIG (default: /var/lib/coda/config_*.json)
  the runtime config generated by the HF packaging
- CREATE_RUNTIME_CONFIG (default: mina-hf-create-runtime-config)
  a copy of ./scripts/hardfork/create_runtime_config.sh
- GENESIS_LEDGER_DIR (default: /var/lib/coda)
  where the tarfiles to verify are stored
- FORKING_FROM_CONFIG_JSON (default: /var/lib/coda/mainnet.json)
  the pre-fork genesis ledger
- SECONDS_PER_SLOT (default: 180)
- PRECOMPUTED_FORK_BLOCK (default: fetches with gsutil)
- GSUTIL (default: gsutil)
  the Google Cloud Storage utility if the PRECOMPUTED_FORK_BLOCK isn't a file
- MINA_LOG_LEVEL (default: info)
  the log level to run the Mina node at when exporting ledgers

Ensures:
- The accounts listed in config.json are the ones in the PACKAGED_DAEMON_CONFIG
- The genesis ledger directory tarfile contents match reference copies created
  fresh from config.json. This takes a long time (>20min) due to rehashing.

Outputs:
- Exit code 0 if validated, 1 otherwise.
EOF

    exit 1
}

PRECOMPUTED_BLOCK_PREFIX=""
CACHED_HARDFORK_DATA=""
CHECKS="config,tarballs,ledgers"

# Track if any errors occurred
error=0

# Parse named flags
while [[ $# -gt 0 ]]; do
    case $1 in
        --network)
            NETWORK_NAME="$2"
            shift 2
            ;;
        --fork-config)
            FORK_CONFIG="$2"
            shift 2
            ;;
        --working-dir)
            WORKDIR="$2"
            shift 2
            ;;
        --precomputed-block-prefix)
            PRECOMPUTED_BLOCK_PREFIX="$2"
            shift 2
            ;;
        --reference-data-dir)
            REFERENCE_DATA_DIR="$2"
            shift 2
            ;;
        --cached-hardfork-data)
            CACHED_HARDFORK_DATA="$2"
            shift 2
            ;;
        --checks)
            CHECKS="$2"
            shift 2
            ;;
        -*)
            echo "Unknown option $1"
            usage
            ;;
        *)
            echo "Unexpected argument $1"
            usage
            ;;
    esac
done


# Validate required arguments
if [[ -z "$NETWORK_NAME" ]]; then
    echo "Error: --network is required"
    usage
fi

if [[ -z "$FORK_CONFIG" ]]; then
    echo "Error: --fork-config is required"
    usage
fi

if [[ -z "$WORKDIR" ]]; then
    echo "Error: --working-dir is required"
    usage
fi

function checks_contains() {
    local check=$1
    IFS=',' read -ra CHECKS_ARRAY <<< "$CHECKS"
    for c in "${CHECKS_ARRAY[@]}"; do
        if [[ "$c" == "$check" ]]; then
            return 0
        fi
    done
    return 1
}

source_build_fallback() {
    v=$(command -v "$1")
    if [ ! -x "$v" ]; then
        if [ ! -x "$2" ]; then
            echo "Error: program not found in PATH (as $1) or relative to cwd (as $2)" >&2
            exit 1
        else
            echo "$2"
        fi
    else
      echo "$1"
    fi
}

MINA_EXE=${MINA_EXE:-$(source_build_fallback mina ./_build/default/src/app/cli/src/mina.exe)}
MINA_GENESIS_EXE=${MINA_GENESIS_EXE:-$(source_build_fallback mina-create-genesis ./_build/default/src/app/runtime_genesis_ledger/runtime_genesis_ledger.exe)}
MINA_LEGACY_GENESIS_EXE=${MINA_LEGACY_GENESIS_EXE:-$(source_build_fallback mina-create-prefork-genesis ./runtime_genesis_ledger_of_mainnet.exe)}
CREATE_RUNTIME_CONFIG=${CREATE_RUNTIME_CONFIG:-$(source_build_fallback mina-hf-create-runtime-config ./scripts/hardfork/create_runtime_config.sh)}
MINA_ROCKSDB_SCANNER=${MINA_ROCKSDB_SCANNER:-$(source_build_fallback mina-rocksdb-scanner ./_build/default/src/app/rocksdb-scanner/rocksdb_scanner.exe)}
GSUTIL=${GSUTIL:-$(command -v gsutil || echo "")}

if [[ ! -e "${PRECOMPUTED_FORK_BLOCK:-}" && ( -z "$GSUTIL" || ! -x "$GSUTIL" ) ]]; then
    echo "Error: gsutil is required when PRECOMPUTED_FORK_BLOCK is nonexistent path"
    exit 1
fi

installed_config=$(echo /var/lib/coda/config_*.json)
PACKAGED_DAEMON_CONFIG=${PACKAGED_DAEMON_CONFIG:-$installed_config}
if [ ! -e "$PACKAGED_DAEMON_CONFIG" ]; then
    echo "Error: set PACKAGED_DAEMON_CONFIG to the path to the JSON file to verify"
    exit 1
fi

GENESIS_LEDGER_DIR=${GENESIS_LEDGER_DIR:-"/var/lib/coda"}

SECONDS_PER_SLOT=${SECONDS_PER_SLOT:=180}

export FORKING_FROM_CONFIG_JSON=${FORKING_FROM_CONFIG_JSON:-$(source_build_fallback /var/lib/coda/mainnet.json genesis_ledgers/mainnet.json)}

export MINA_LIBP2P_PASS=''

mkdir -p "$WORKDIR"/{ledgers{,-backup},keys}
chmod 700 "$WORKDIR/keys"

fork_block_state_hash=$(jq -r '.proof.fork.state_hash' "$FORK_CONFIG")
fork_block_length=$(jq -r '.proof.fork.blockchain_length' "$FORK_CONFIG")

# Fetch the precomputed fork block from local file or gcloud
fetch_precomputed_fork_block() {
    local network_name=$1
    local fork_block_length=$2
    local fork_block_state_hash=$3
    local precomputed_block_prefix=$4
    local dst=$5

    # Put the fork block where we want it, fetch it from gcloud if necessary
    if [[ -z "${PRECOMPUTED_FORK_BLOCK:-}" ]]; then
        PRECOMPUTED_FORK_BLOCK="${precomputed_block_prefix:-gs://mina_network_block_data/$network_name}-$fork_block_length-$fork_block_state_hash.json"
    fi

    local src="$PRECOMPUTED_FORK_BLOCK"
    if [[ -e "$src" ]]; then
        cp "$src" "$dst"
    else
        "$GSUTIL" cp "$src" "$dst"
    fi
}

# export the packaged ledgers in a way where we know which one is which
function extract_ledgers() {
    config_file=$1
    ledger_dir=$2
    json_prefix=$3

    "$MINA_EXE" daemon --libp2p-keypair "$WORKDIR/keys/p2p" \
      --config-file "$config_file" --config-file "$override_file" \
      --seed --genesis-ledger-dir "$ledger_dir" \
      --log-level "${MINA_LOG_LEVEL:-info}" &

    while ! "$MINA_EXE" ledger export staged-ledger  | jq . 2>/dev/null >"$json_prefix-staged.json"; do
        sleep 1m
        if ! grep -qFx "$(cat ~/.mina-config/.mina-lock)" <(jobs -rp); then
            echo "daemon died before exporting ledgers" >&2
            exit 1
        fi
    done

    "$MINA_EXE" ledger export staking-epoch-ledger | jq . > "$json_prefix-staking.json"
    "$MINA_EXE" ledger export next-epoch-ledger | jq . > "$json_prefix-next.json"

    "$MINA_EXE" client stop
}

already_restored=0
function restore_or_generate_hardfork_data() {
    if [ $already_restored -eq 1 ]; then
        return
    fi
    already_restored=1

    if [ -n "$CACHED_HARDFORK_DATA" ]; then
        echo "Restoring cached hardfork data from $CACHED_HARDFORK_DATA" >&2

        # Check if cached files exist
        if [ ! -d "$CACHED_HARDFORK_DATA/ledgers" ]; then
            echo "Warning: Cached ledgers directory not found at $CACHED_HARDFORK_DATA/ledgers" >&2
            echo "Falling back to generating ledgers..." >&2
            "$MINA_GENESIS_EXE" --config-file "$FORK_CONFIG" --genesis-dir "$WORKDIR/ledgers" --hash-output-file "$WORKDIR/hashes.json"
        elif [ ! -f "$CACHED_HARDFORK_DATA/hashes.json" ]; then
            echo "Warning: Cached hashes.json file not found at $CACHED_HARDFORK_DATA/hashes.json" >&2
            echo "Falling back to generating ledgers..." >&2
            "$MINA_GENESIS_EXE" --config-file "$FORK_CONFIG" --genesis-dir "$WORKDIR/ledgers" --hash-output-file "$WORKDIR/hashes.json"
        else
            # Copy cached files
            cp -r "$CACHED_HARDFORK_DATA"/ledgers/*.tar.gz "$WORKDIR/ledgers"/
            cp "$CACHED_HARDFORK_DATA/hashes.json" "$WORKDIR/hashes.json"
            echo "Successfully restored cached hardfork data" >&2
        fi
    else
        "$MINA_GENESIS_EXE" --config-file "$FORK_CONFIG" --genesis-dir "$WORKDIR/ledgers" --hash-output-file "$WORKDIR/hashes.json"
    fi
}

if [ -z "${REFERENCE_DATA_DIR:-}" ]; then
    echo "generating genesis ledgers ... (this may take a while)" >&2
    "$MINA_LEGACY_GENESIS_EXE" --config-file "$FORK_CONFIG" --genesis-dir "$WORKDIR/legacy_ledgers" --hash-output-file "$WORKDIR/legacy_hashes.json"
    REFERENCE_DATA_DIR="$WORKDIR"
else
    if [[ ! -d "$REFERENCE_DATA_DIR" ]]; then
        echo "Error: reference data dir $REFERENCE_DATA_DIR does not exist" >&2
        exit 1
    fi
    echo "using reference data dir: $REFERENCE_DATA_DIR. Skipping legacy ledger generation." >&2
fi

if checks_contains config; then
    echo "Validating config..." >&2

    fetch_precomputed_fork_block "$NETWORK_NAME" "$fork_block_length" "$fork_block_state_hash" "$PRECOMPUTED_BLOCK_PREFIX" "$WORKDIR/precomputed_fork_block.json"

    result=$(jq --slurpfile block "$WORKDIR/precomputed_fork_block.json" \
    --slurpfile legacy_hashes "$REFERENCE_DATA_DIR/legacy_hashes.json" -n '
    ($legacy_hashes[0].epoch_data.staking.hash == $block[0].data.protocol_state.body.consensus_state.staking_epoch_data.ledger.hash and
    $legacy_hashes[0].epoch_data.next.hash == $block[0].data.protocol_state.body.consensus_state.next_epoch_data.ledger.hash and
    $legacy_hashes[0].ledger.hash == $block[0].data.protocol_state.body.blockchain_state.staged_ledger_hash.non_snark.ledger_hash)')

    if [ "$result" != "true" ]; then
        echo "Hashes in config $FORK_CONFIG don't match hashes from the precomputed block $PRECOMPUTED_FORK_BLOCK" >&2
        exit 1
    fi
else
    echo "Skipping config validation as per --checks option" >&2
fi

if checks_contains ledgers; then
    restore_or_generate_hardfork_data

    FORK_CONFIG_JSON="$FORK_CONFIG" \
    LEDGER_HASHES_JSON="$WORKDIR/hashes.json" \
    GENESIS_TIMESTAMP=$(jq -r '.genesis.genesis_state_timestamp' "$PACKAGED_DAEMON_CONFIG") \
    "$CREATE_RUNTIME_CONFIG" > "$WORKDIR/config-substituted.json"

    echo "exporting ledgers from running node ... (this may take a while)" >&2

    if [ ! -e "$WORKDIR/keys/p2p" ]; then
        "$MINA_EXE" libp2p generate-keypair --privkey-path "$WORKDIR/keys/p2p"
    fi

    override_file="$WORKDIR/override-genesis-timestamp.json"
    echo "{\"genesis\":{\"genesis_state_timestamp\":\"$(date -u +"%Y-%m-%dT%H:%M:%SZ")\"}}" > "$override_file"

    mkdir -p "$WORKDIR/ledgers-downloaded"

    mv -t "$WORKDIR/ledgers-backup" "$GENESIS_LEDGER_DIR"/*.tar.gz
    if [ -z ${NO_TEST_LEDGER_DOWNLOAD+z} ]; then
        # We removed all *.tar.gz from GENESIS_LEDGER_DIR (in case it's one of default paths)
        # and expect Mina node to be able to download these from S3
        extract_ledgers "$PACKAGED_DAEMON_CONFIG" "$WORKDIR/ledgers-downloaded" "$WORKDIR/downloaded"
        rm -Rf /tmp/coda_cache_dir/*.tar.gz /tmp/s3_cache_dir/*.tar.gz
    fi
    extract_ledgers "$WORKDIR/config-substituted.json" "$WORKDIR/ledgers" "$WORKDIR/reference"
    mv -t "$GENESIS_LEDGER_DIR" "$WORKDIR/ledgers-backup"/*
    extract_ledgers "$PACKAGED_DAEMON_CONFIG" "$GENESIS_LEDGER_DIR" "$WORKDIR/packaged"

    echo "Performing final comparisons..." >&2

    result=$(jq --slurpfile a "$WORKDIR/config-substituted.json" --slurpfile b "$PACKAGED_DAEMON_CONFIG" -n '
    ($a[0].epoch_data.staking.hash == $b[0].epoch_data.staking.hash and
    $a[0].epoch_data.next.hash == $b[0].epoch_data.next.hash and
    $a[0].ledger.hash == $b[0].ledger.hash)')

    if [ "$result" != "true" ]; then
        echo "Packaged config hashes in $PACKAGED_DAEMON_CONFIG not expected compared to $WORKDIR/config-substituted.json" >&2
        error=1
    fi

    for file in "$WORKDIR"/packaged-*.json; do
        name=$(basename "$file")
        name=${name%.json}
        name=${name#packaged-}

        if ! cmp "$file" "$WORKDIR/reference-$name.json"; then
            echo "Error: $file does not match reference" >&2
            error=1
        fi
        if [ -z ${NO_TEST_LEDGER_DOWNLOAD+z} ]; then
            if ! cmp "$file" "$WORKDIR/downloaded-$name.json"; then
                echo "Error: $file does not match downloaded" >&2
                error=1
            fi
        fi
    done
else
    echo "Skipping ledgers validation as per --checks option" >&2
fi

if checks_contains tarballs; then

    restore_or_generate_hardfork_data

    echo "Validating tarballs..." >&2
    for file in "$WORKDIR"/ledgers/*.tar.gz; do
        tarname=$(basename "$file")
        tarname=${tarname%.tar.gz}
        tardir="$WORKDIR/ledgers/$tarname"
        mkdir -p "$tardir"/{packaged,generated,web}
        tar -xzf "$file" -C "$tardir/generated"
        tar -xzf "$GENESIS_LEDGER_DIR/$tarname.tar.gz" -C "$tardir/packaged"

        base_s3_url="${MINA_LEDGER_S3_BUCKET:-https://s3-us-west-2.amazonaws.com/snark-keys-ro.o1test.net}"
        curl "$base_s3_url/$tarname.tar.gz" | tar -xz -C "$tardir/web"

        declare -A scanner_pids
        for kind in packaged web generated; do
            $MINA_ROCKSDB_SCANNER dump --db-path "$tardir/$kind" --output-file "$WORKDIR/$kind.scan" &
            scanner_pids[$kind]=$!
        done

        # Wait for all scanner jobs to finish
        for kind in packaged web generated; do
            wait ${scanner_pids[$kind]}
        done

        if ! ( cmp -s "$WORKDIR/generated.scan" "$WORKDIR/packaged.scan" && \
            cmp -s "$WORKDIR/packaged.scan" "$WORKDIR/web.scan" ); then
            echo "Error: kvdb contents mismatch for $tarname" >&2
            error=1
        fi
    done
else
    echo "Skipping tarballs validation as per --checks option" >&2
fi

if [ $error -ne 0 ]; then
    echo "Error: failed validation" >&2
    exit 1
else
    echo "Validation successful" >&2
fi
