#!/usr/bin/env bash

set -euo pipefail

parse_size() {
    local input="$1"
    local number suffix

    if [[ "$input" =~ ^([0-9]+)([KkMmGg]?)$ ]]; then
    number="${BASH_REMATCH[1]}"
    suffix="${BASH_REMATCH[2]}"

    case "$suffix" in
    K|k) echo $((number * 1024)) ;;
    M|m) echo $((number * 1024 * 1024)) ;;
    G|g) echo $((number * 1024 * 1024 * 1024)) ;;
    "")  echo "$number" ;;
    *)   return 1 ;;
    esac
    else
    echo "Invalid size: $input" >&2
    return 1
    fi
}

cleanup() {
    rm -f buffer.json descriptor.bin host_recv.stdin
}
trap cleanup EXIT INT TERM

if [[ $# -lt 4 || $# -gt 5 ]]; then
    echo "Usage: $0 <tag> <runs> <bytes|size> <remote-input-file> [remote-log-prefix-base]"
    echo "Example: $0 VANILLA 10 256M bins/256M.bin"
    exit 1
fi

TAG="$1"
RUNS="$2"
DEBUG_BYTES="$(parse_size "$3")"
REMOTE_INPUT="$4"
REMOTE_PREFIX_BASE="${TAG}_$3"

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

DEBUG_OUT="./test.bin"

# Host-side PCI BDF toward DPU.
PCI_TO_DPU="0000:5e:00.0"

# 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
)

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

echo "====="
echo "SIHO'S DMA BENCHMARKING"
echo "SSH KEY LOCATION: $SSH_KEY"
echo "REMOTE HOST: $REMOTE_HOST"
echo "REMOTE DIR: $REMOTE_DIR"
echo "TAG: $TAG"
echo "HOST PCIE ADDRESS TO DPU: $PCI_TO_DPU"
echo "DPU PCIE ADDRESS: $DPU_PCI"
echo "RUNS: $RUNS"
echo "BYTES: $DEBUG_BYTES"
echo "REMOTE INPUT: $REMOTE_INPUT"
echo "====="
echo "PROCEED (y/N)?"

read -r proceed

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

sudo -v

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

    rm -f descriptor.bin buffer.json host_recv.stdin "$DEBUG_OUT"
    mkfifo host_recv.stdin

    # Open FIFO read-write so the receiver does not get accidental EOF.
    exec 3<>host_recv.stdin

    sudo -n ./doca_dma_host_recv \
	        -p "$PCI_TO_DPU" \
		    -o "$DEBUG_OUT" \
		        -s "$DEBUG_BYTES" \
			    < host_recv.stdin &
    host_pid=$!

    ok=0
    for _ in $(seq 1 100); do
	    if [[ -s descriptor.bin && -s buffer.json ]]; then
		    ok=1
		    break
	    fi

	    sleep 0.1
    done

    if [[ "$ok" != "1" ]]; then
	    echo "ERROR: descriptor.bin or buffer.json was not generated within timeout"
	    echo "host_pid=$host_pid"

	    if kill -0 "$host_pid" 2>/dev/null; then
		    echo "host receiver still appears to be running"
	    else
		    echo "host receiver process is not visible anymore"
	    fi

	    kill "$host_pid" 2>/dev/null || true
	    exec 3>&- 2>/dev/null || true
	    rm -f host_recv.stdin
	    exit 1
    fi

    echo "[host] Copying descriptor.bin and buffer.json to DPU..."
    if ! scp "${SCP_OPTS[@]}" descriptor.bin buffer.json "$REMOTE_HOST:$REMOTE_DIR/"; then
        echo "ERROR: scp failed"
        kill "$host_pid" 2>/dev/null || true
        exec 3>&-
        exit 1
    fi

    echo "[dpu] Running DMA sender through wrapper..."
    if ! dpu_out="$(
    ssh "${SSH_OPTS[@]}" "$REMOTE_HOST" \
    "cd '$REMOTE_DIR' && sudo -n '$DPU_SCRIPT' '$REMOTE_PREFIX_BASE' '$DPU_PCI' '$REMOTE_INPUT'" \
    2>&1
    )"; then
    echo "ERROR: DPU sender failed"
    echo "$dpu_out"
    kill "$host_pid" 2>/dev/null || true
    exec 3>&-
    exit 1
    fi

    bytes="$(awk '/bytes:/ {print $2}' <<< "$dpu_out" | tail -n1)"
    time_val="$(awk '/time:/ {print $2}' <<< "$dpu_out" | tail -n1)"
    time_unit="$(awk '/time:/ {print $3}' <<< "$dpu_out" | tail -n1)"
    bw_val="$(awk '/BW:/ {print $2}' <<< "$dpu_out" | tail -n1)"
    bw_unit="$(awk '/BW:/ {print $3}' <<< "$dpu_out" | tail -n1)"

    if [[ -z "${bytes:-}" || -z "${time_val:-}" || -z "${bw_val:-}" ]]; then
        echo "ERROR: failed to parse DPU output"
        echo "$dpu_out"
        kill "$host_pid" 2>/dev/null || true
        exec 3>&-
        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 "[host] Releasing host receiver..."
    printf '\n' >&3
    exec 3>&-

    if ! wait "$host_pid"; then
        echo "ERROR: host receiver failed"
        exit 1
    fi

    rm -f host_recv.stdin

    echo "[result] bytes=$bytes time_usec=$time_usec bw_KiBps=$bw_kibps"
    echo "===== RUN $i/$RUNS DONE ====="
done

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