#!/usr/bin/env bash

set -euo pipefail

usage() {
    echo "Usage: $0 <tag> <runs> <remote-chunk-directory> <chunk-size> [remote-log-prefix-base]"
    echo
    echo "Examples:"
    echo "  $0 VANILLA_1G 10 bins/1G 1M"
    echo "  $0 TRACE_5G   10 bins/5G 1M TRACE_5G"
}

if [[ $# -lt 4 || $# -gt 5 ]]; then
    usage
    exit 1
fi

TAG="$1"
RUNS="$2"
REMOTE_CHUNK_DIR="$3"
CHUNK_SIZE="$4"
REMOTE_PREFIX_BASE="${5:-${TAG}}"

SSH_KEY="$HOME/.ssh/id_ed25519"
REMOTE_HOST="ubuntu@143.248.53.46"
REMOTE_DIR="/home/ubuntu/siho-benchmark/compress"

# DPU-side PCI BDF.
DPU_PCI="03:00.0"

# DPU-side wrapper script.
DPU_SCRIPT="./run.sh"

TIMES_CSV="times_${TAG}.csv"
BW_CSV="bw_${TAG}.csv"

SSH_OPTS=(
    -i "$SSH_KEY"
    -n
    -o BatchMode=yes
    -o StrictHostKeyChecking=accept-new
)

echo "====="
echo "SIHO'S DOCA DECOMPRESS BENCHMARKING"
echo "SSH KEY LOCATION: $SSH_KEY"
echo "REMOTE HOST: $REMOTE_HOST"
echo "REMOTE DIR: $REMOTE_DIR"
echo "TAG: $TAG"
echo "DPU PCIE ADDRESS: $DPU_PCI"
echo "RUNS: $RUNS"
echo "REMOTE CHUNK DIR: $REMOTE_CHUNK_DIR"
echo "CHUNK SIZE: $CHUNK_SIZE"
echo "DPU SCRIPT: $DPU_SCRIPT"
echo "====="
echo "PROCEED (y/N)?"

read -r proceed

if [[ "${proceed:-}" != "y" ]]; then
    echo "TERMINATING..."
    exit 0
fi

for i in $(seq 1 "$RUNS"); do
    echo "===== RUN $i/$RUNS ====="

    RUN_LABEL="${REMOTE_PREFIX_BASE}_run_${i}"

    echo "[dpu] Running decompressor through wrapper..."

    if ! dpu_out="$(
        ssh "${SSH_OPTS[@]}" "$REMOTE_HOST" \
            "cd '$REMOTE_DIR' && sudo -n '$DPU_SCRIPT' '$REMOTE_PREFIX_BASE' '$DPU_PCI' '$REMOTE_CHUNK_DIR' '$CHUNK_SIZE' '$RUN_LABEL'" \
            2>&1
    )"; then
        echo "ERROR: DPU decompressor failed"
        echo "$dpu_out"
        exit 1
    fi

    compressed_bytes="$(
        awk '/^[[:space:]]*compressed_bytes:/ {print $2}' <<< "$dpu_out" | tail -n1
    )"

    bytes="$(
        awk '/^[[:space:]]*bytes:/ {print $2}' <<< "$dpu_out" | tail -n1
    )"

    time_val="$(
        awk '/^[[:space:]]*time:/ {print $2}' <<< "$dpu_out" | tail -n1
    )"

    time_unit="$(
        awk '/^[[:space:]]*time:/ {print $3}' <<< "$dpu_out" | tail -n1
    )"

    bw_val="$(
        awk '/^[[:space:]]*BW:/ {print $2}' <<< "$dpu_out" | tail -n1
    )"

    bw_unit="$(
        awk '/^[[:space:]]*BW:/ {print $3}' <<< "$dpu_out" | tail -n1
    )"

    chunks="$(
        awk '/^[[:space:]]*chunks:/ {print $2}' <<< "$dpu_out" | tail -n1
    )"

    if [[ -z "${bytes:-}" || -z "${time_val:-}" || -z "${bw_val:-}" ]]; then
        echo "ERROR: failed to parse DPU output"
        echo "$dpu_out"
        exit 1
    fi

    # Convert time to usec.
    case "$time_unit" in
    usec) time_usec="$time_val" ;;
    msec) time_usec="$(awk -v x="$time_val" 'BEGIN { printf "%.6f", x * 1000.0 }')" ;;
    sec)  time_usec="$(awk -v x="$time_val" 'BEGIN { printf "%.6f", x * 1000000.0 }')" ;;
    nsec) time_usec="$(awk -v x="$time_val" 'BEGIN { printf "%.6f", x / 1000.0 }')" ;;
    *)    time_usec="$time_val" ;;
    esac

    # Convert bandwidth to KiB/s.
    case "$bw_unit" in
    KiB/s|KiBps) bw_kibps="$bw_val" ;;
    MiB/s|MiBps) bw_kibps="$(awk -v x="$bw_val" 'BEGIN { printf "%.6f", x * 1024.0 }')" ;;
    GiB/s|GiBps) bw_kibps="$(awk -v x="$bw_val" 'BEGIN { printf "%.6f", x * 1024.0 * 1024.0 }')" ;;
    B/s|Bps)     bw_kibps="$(awk -v x="$bw_val" 'BEGIN { printf "%.6f", x / 1024.0 }')" ;;
    *)           bw_kibps="$bw_val" ;;
    esac

    if [[ ! -f "$TIMES_CSV" ]]; then
        echo "bytes,time_usec" > "$TIMES_CSV"
    fi

    if [[ ! -f "$BW_CSV" ]]; then
        echo "bytes,bw_KiBps" > "$BW_CSV"
    fi

    echo "$bytes,$time_usec" >> "$TIMES_CSV"
    echo "$bytes,$bw_kibps" >> "$BW_CSV"

    echo "[result] chunks=${chunks:-N/A} compressed_bytes=${compressed_bytes:-N/A} bytes=$bytes time_usec=$time_usec bw_KiBps=$bw_kibps"
    echo "===== RUN $i/$RUNS DONE ====="
done

echo "Wrote: $TIMES_CSV"
echo "Wrote: $BW_CSV"