#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"

require_cmd() {
  command -v "$1" >/dev/null 2>&1 || {
    echo "missing required command: $1" >&2
    exit 1
  }
}

for cmd in curl qemu-io; do
  require_cmd "${cmd}"
done

if ! command -v jq >/dev/null 2>&1 && ! command -v python3 >/dev/null 2>&1; then
  echo "missing required command: jq or python3" >&2
  exit 1
fi

json_get() {
  local query="$1"
  if command -v jq >/dev/null 2>&1; then
    jq -r "${query}"
  else
    python3 -c 'import json,sys
data=json.load(sys.stdin)
value=data
for part in sys.argv[1].split("."):
    if not part:
        continue
    value=value.get(part) if isinstance(value, dict) else None
    if value is None:
        break
print("" if value is None else value)
' "${query}"
  fi
}

RUN_ID="${CORONAFS_BENCH_RUN_ID:-$$}"
LISTEN_PORT="${CORONAFS_BENCH_PORT:-$((25088 + (RUN_ID % 1000)))}"
EXPORT_BASE_PORT="${CORONAFS_BENCH_EXPORT_BASE_PORT:-$((26100 + (RUN_ID % 1000)))}"
VOLUME_ID="${CORONAFS_BENCH_VOLUME_ID:-local-bench-${RUN_ID}}"
SIZE_MIB="${CORONAFS_BENCH_SIZE_MIB:-${CORONAFS_BENCH_SIZE_MB:-512}}"
SIZE_BYTES="${CORONAFS_BENCH_SIZE_BYTES:-$((SIZE_MIB * 1024 * 1024))}"
WORKLOAD_MIB="${CORONAFS_BENCH_WORKLOAD_MIB:-${CORONAFS_BENCH_WORKLOAD_MB:-256}}"
EXPORT_CACHE_MODE="${CORONAFS_BENCH_EXPORT_CACHE_MODE:-none}"
EXPORT_AIO_MODE="${CORONAFS_BENCH_EXPORT_AIO_MODE:-threads}"
EXPORT_DISCARD_MODE="${CORONAFS_BENCH_EXPORT_DISCARD_MODE:-ignore}"
EXPORT_DETECT_ZEROES_MODE="${CORONAFS_BENCH_EXPORT_DETECT_ZEROES_MODE:-off}"
SERVER_BIN="${CORONAFS_SERVER_BIN:-}"

if (( WORKLOAD_MIB > SIZE_MIB )); then
  echo "workload ${WORKLOAD_MIB} MiB exceeds volume size ${SIZE_MIB} MiB" >&2
  exit 1
fi

if [[ -z "${SERVER_BIN}" ]]; then
  SERVER_CMD=(
    cargo run
    --manifest-path "${REPO_ROOT}/coronafs/Cargo.toml"
    -p coronafs-server
    --
  )
else
  SERVER_CMD=("${SERVER_BIN}")
fi

TMP_DIR="$(mktemp -d)"
CONFIG_PATH="${TMP_DIR}/coronafs.toml"
SERVER_LOG="${TMP_DIR}/coronafs.log"
SERVER_PID=""

show_server_log() {
  if [[ -f "${SERVER_LOG}" ]]; then
    echo "--- coronafs server log ---" >&2
    tail -n 200 "${SERVER_LOG}" >&2 || true
    echo "--- end coronafs server log ---" >&2
  fi
}

delete_volume_if_present() {
  curl -fsS -X DELETE "http://127.0.0.1:${LISTEN_PORT}/v1/volumes/${VOLUME_ID}" >/dev/null 2>&1 || true
}

cleanup() {
  delete_volume_if_present
  local pid_file="${TMP_DIR}/data/pids/${VOLUME_ID}.pid"
  if [[ -f "${pid_file}" ]]; then
    local export_pid=""
    export_pid="$(tr -d '\n' <"${pid_file}" 2>/dev/null || true)"
    if [[ -n "${export_pid}" ]] && kill -0 "${export_pid}" 2>/dev/null; then
      kill "${export_pid}" >/dev/null 2>&1 || true
      wait "${export_pid}" >/dev/null 2>&1 || true
    fi
    rm -f "${pid_file}"
  fi
  if [[ -n "${SERVER_PID}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then
    kill "${SERVER_PID}" >/dev/null 2>&1 || true
    wait "${SERVER_PID}" >/dev/null 2>&1 || true
  fi
  rm -rf "${TMP_DIR}"
}
trap cleanup EXIT

cat >"${CONFIG_PATH}" <<EOF
listen_addr = "127.0.0.1:${LISTEN_PORT}"
advertise_host = "127.0.0.1"
data_dir = "${TMP_DIR}/data"
export_bind_addr = "127.0.0.1"
export_base_port = ${EXPORT_BASE_PORT}
export_port_count = 8
export_shared_clients = 32
export_cache_mode = "${EXPORT_CACHE_MODE}"
export_aio_mode = "${EXPORT_AIO_MODE}"
export_discard_mode = "${EXPORT_DISCARD_MODE}"
export_detect_zeroes_mode = "${EXPORT_DETECT_ZEROES_MODE}"
preallocate = false
sync_on_write = false
log_level = "info"
EOF

"${SERVER_CMD[@]}" --config "${CONFIG_PATH}" >"${SERVER_LOG}" 2>&1 &
SERVER_PID="$!"

deadline=$((SECONDS + 60))
until curl -fsS "http://127.0.0.1:${LISTEN_PORT}/healthz" >/dev/null 2>&1; do
  if (( SECONDS >= deadline )); then
    echo "timed out waiting for coronafs local bench server" >&2
    tail -n 200 "${SERVER_LOG}" >&2 || true
    exit 1
  fi
  sleep 1
done

create_response_file="${TMP_DIR}/create-response.txt"
create_status="$(
  curl -sS \
    -o "${create_response_file}" \
    -w '%{http_code}' \
    -X PUT \
    -H 'content-type: application/json' \
    -d "{\"size_bytes\":${SIZE_BYTES}}" \
    "http://127.0.0.1:${LISTEN_PORT}/v1/volumes/${VOLUME_ID}"
)"
if [[ "${create_status}" -lt 200 || "${create_status}" -ge 300 ]]; then
  echo "failed to create CoronaFS benchmark volume: HTTP ${create_status}" >&2
  cat "${create_response_file}" >&2 || true
  show_server_log
  exit 1
fi

export_response_file="${TMP_DIR}/export-response.txt"
export_status="$(
  curl -sS \
    -o "${export_response_file}" \
    -w '%{http_code}' \
    -X POST \
    "http://127.0.0.1:${LISTEN_PORT}/v1/volumes/${VOLUME_ID}/export"
)"
if [[ "${export_status}" -lt 200 || "${export_status}" -ge 300 ]]; then
  echo "failed to export CoronaFS benchmark volume: HTTP ${export_status}" >&2
  cat "${export_response_file}" >&2 || true
  show_server_log
  exit 1
fi
EXPORT_JSON="$(cat "${export_response_file}")"
EXPORT_URI="$(printf '%s' "${EXPORT_JSON}" | json_get '.export.uri')"
[[ -n "${EXPORT_URI}" && "${EXPORT_URI}" != "null" ]] || {
  echo "failed to obtain CoronaFS export URI" >&2
  printf '%s\n' "${EXPORT_JSON}" >&2
  show_server_log
  exit 1
}

run_qemu_io() {
  local extra=()
  local start_ns end_ns elapsed_ns
  local args=("$@")
  local cmd=()
  local qemu_cmd=""

  if [[ "${#args[@]}" -eq 0 ]]; then
    echo "run_qemu_io requires at least one qemu-io command" >&2
    exit 1
  fi

  while [[ "${#args[@]}" -gt 0 && "${args[0]}" == --* ]]; do
    extra+=("${args[0]}")
    args=("${args[@]:1}")
  done

  cmd=(qemu-io -f raw "${extra[@]}")
  for qemu_cmd in "${args[@]}"; do
    cmd+=(-c "${qemu_cmd}")
  done
  cmd+=("${EXPORT_URI}")

  start_ns="$(date +%s%N)"
  "${cmd[@]}" >/dev/null
  end_ns="$(date +%s%N)"
  elapsed_ns="$((end_ns - start_ns))"
  printf '%s\n' "${elapsed_ns}"
}

calc_mib_per_s() {
  local bytes="$1"
  local elapsed_ns="$2"
  awk -v bytes="${bytes}" -v elapsed_ns="${elapsed_ns}" '
    BEGIN {
      if (elapsed_ns <= 0) {
        print "0.00"
      } else {
        printf "%.2f", (bytes / 1048576.0) / (elapsed_ns / 1000000000.0)
      }
    }
  '
}

BYTES="$((WORKLOAD_MIB * 1024 * 1024))"
WRITE_NS="$(run_qemu_io "write -P 0x5a 0 ${WORKLOAD_MIB}M" "flush")"
READ_NS="$(run_qemu_io "read -P 0x5a 0 ${WORKLOAD_MIB}M")"
WRITE_MIBPS="$(calc_mib_per_s "${BYTES}" "${WRITE_NS}")"
READ_MIBPS="$(calc_mib_per_s "${BYTES}" "${READ_NS}")"

printf 'CoronaFS local export bench: uri=%s cache=%s aio=%s write=%s MiB/s read=%s MiB/s size=%s MiB\n' \
  "${EXPORT_URI}" "${EXPORT_CACHE_MODE}" "${EXPORT_AIO_MODE}" "${WRITE_MIBPS}" "${READ_MIBPS}" "${WORKLOAD_MIB}"

printf '%s\t%s\t%s\t%s\t%s\n' "${EXPORT_URI}" "${EXPORT_CACHE_MODE}" "${EXPORT_AIO_MODE}" "${WRITE_MIBPS}" "${READ_MIBPS}"