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

# M5: Generate perf profiles + flamegraphs + hotspots for Stage2 objects.
# Default target is ROA max sample since it is currently the largest gap in
# ours-vs-routinator decode comparison.
#
# Outputs (under rpki/target/bench):
# - perf_stage2_ours_<type>_<sample>_decode.data
# - flamegraph_stage2_ours_<type>_<sample>_decode_release.svg
# - hotspots_stage2_ours_<type>_<sample>_decode_release.txt
# - perf_stage2_ours_<type>_<sample>_landing.data
# - flamegraph_stage2_ours_<type>_<sample>_landing_release.svg
# - hotspots_stage2_ours_<type>_<sample>_landing_release.txt
# - perf_stage2_routinator_<type>_<sample>_decode.data
# - flamegraph_stage2_routinator_<type>_<sample>_decode_release.svg
# - hotspots_stage2_routinator_<type>_<sample>_decode_release.txt

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
OUT_DIR="$ROOT_DIR/target/bench"
mkdir -p "$OUT_DIR"

TYPE="${BENCH_TYPE:-roa}"       # cer|crl|roa
SAMPLE="${BENCH_SAMPLE:-max}"   # e.g. max|p99|p50

ITERS="${BENCH_ITERS:-2000}"
WARMUP="${BENCH_WARMUP_ITERS:-50}"
FREQ="${FLAMEGRAPH_FREQ:-99}"

PROFILE_LANDING="${PROFILE_LANDING:-1}" # 0|1

if ! command -v perf >/dev/null 2>&1; then
  echo "ERROR: perf not found." >&2
  exit 2
fi
if ! command -v flamegraph >/dev/null 2>&1; then
  echo "ERROR: flamegraph not found. Install with: cargo install flamegraph" >&2
  exit 2
fi

# On WSL2, /usr/bin/perf is often a wrapper that errors because there is no
# kernel-matched perf binary. In that case, prefer a real perf binary under
# /usr/lib/linux-tools/*/perf by putting a shim earlier in PATH.
PERF_WRAPPER_OK=1
PERF_VERSION_OUT="$(perf --version 2>&1 || true)"
if echo "${PERF_VERSION_OUT}" | grep -q "WARNING: perf not found for kernel"; then
  PERF_WRAPPER_OK=0
fi

if [[ "${PERF_WRAPPER_OK}" == "0" ]]; then
  PERF_REAL="$(ls -1 /usr/lib/linux-tools/*/perf 2>/dev/null | head -n 1 || true)"
  if [[ -z "${PERF_REAL}" ]]; then
    echo "ERROR: perf wrapper found, but no real perf binary under /usr/lib/linux-tools/*/perf" >&2
    exit 2
  fi

  SHIM_DIR="$OUT_DIR/tools"
  mkdir -p "${SHIM_DIR}"
  cat > "${SHIM_DIR}/perf" <<EOF
#!/usr/bin/env bash
exec "${PERF_REAL}" "\$@"
EOF
  chmod +x "${SHIM_DIR}/perf"
  export PATH="${SHIM_DIR}:${PATH}"
  echo "Using perf shim -> ${PERF_REAL}" >&2
fi

echo "type=${TYPE} sample=${SAMPLE} iters=${ITERS} warmup=${WARMUP} freq=${FREQ}" >&2

build_ours_test_bin() {
  (cd "$ROOT_DIR" && cargo test --release --test bench_stage2_decode_profile_selected_der_v2 --no-run -q)
  local bin
  bin="$(find "$ROOT_DIR/target/release/deps" -maxdepth 1 -type f -executable -name 'bench_stage2_decode_profile_selected_der_v2-*' | head -n 1 || true)"
  if [[ -z "${bin}" ]]; then
    echo "ERROR: failed to locate ours test binary under target/release/deps" >&2
    exit 2
  fi
  echo "${bin}"
}

build_routinator_bin() {
  (cd "$ROOT_DIR/benchmark/routinator_object_bench" && cargo build --release -q)
  local bin="$ROOT_DIR/benchmark/routinator_object_bench/target/release/routinator-object-bench"
  if [[ ! -x "${bin}" ]]; then
    echo "ERROR: failed to locate routinator-object-bench binary: ${bin}" >&2
    exit 2
  fi
  echo "${bin}"
}

run_perf_profile() {
  local perfdata="$1"
  shift
  echo "[perf record] $perfdata" >&2
  perf record -o "$perfdata" -F "$FREQ" --call-graph dwarf -g -- "$@"
}

gen_flamegraph() {
  local perfdata="$1"
  local out_svg="$2"
  local title="$3"
  local subtitle="$4"
  echo "[flamegraph] $out_svg" >&2
  flamegraph --perfdata "$perfdata" --output "$out_svg" --title "$title" --subtitle "$subtitle" >/dev/null
}

gen_hotspots() {
  local perfdata="$1"
  local out_txt="$2"
  echo "[perf report] $out_txt" >&2
  # Keep it deterministic enough for diffing; avoid call-graph children view.
  perf report --stdio -i "$perfdata" --no-children --sort overhead,symbol,dso > "$out_txt" || true
}

OURS_BIN="$(build_ours_test_bin)"
ROUT_BIN="$(build_routinator_bin)"

# OURS decode+validate
OURS_DECODE_PERF="$OUT_DIR/perf_stage2_ours_${TYPE}_${SAMPLE}_decode.data"
OURS_DECODE_SVG="$OUT_DIR/flamegraph_stage2_ours_${TYPE}_${SAMPLE}_decode_release.svg"
OURS_DECODE_HOT="$OUT_DIR/hotspots_stage2_ours_${TYPE}_${SAMPLE}_decode_release.txt"

run_perf_profile "$OURS_DECODE_PERF" \
  env \
    BENCH_MODE=decode_validate \
    BENCH_TYPE="$TYPE" \
    BENCH_SAMPLE="$SAMPLE" \
    BENCH_ITERS="$ITERS" \
    BENCH_ROUNDS=1 \
    BENCH_WARMUP_ITERS="$WARMUP" \
    "$OURS_BIN" --ignored --nocapture >/dev/null

gen_flamegraph "$OURS_DECODE_PERF" "$OURS_DECODE_SVG" \
  "stage2 ours decode_validate ${TYPE}.${SAMPLE}" "iters=${ITERS} warmup=${WARMUP} freq=${FREQ}"
gen_hotspots "$OURS_DECODE_PERF" "$OURS_DECODE_HOT"

# OURS landing
if [[ "${PROFILE_LANDING}" == "1" ]]; then
  OURS_LAND_PERF="$OUT_DIR/perf_stage2_ours_${TYPE}_${SAMPLE}_landing.data"
  OURS_LAND_SVG="$OUT_DIR/flamegraph_stage2_ours_${TYPE}_${SAMPLE}_landing_release.svg"
  OURS_LAND_HOT="$OUT_DIR/hotspots_stage2_ours_${TYPE}_${SAMPLE}_landing_release.txt"

  run_perf_profile "$OURS_LAND_PERF" \
    env \
      BENCH_MODE=landing \
      BENCH_TYPE="$TYPE" \
      BENCH_SAMPLE="$SAMPLE" \
      BENCH_ITERS="$ITERS" \
      BENCH_ROUNDS=1 \
      BENCH_WARMUP_ITERS="$WARMUP" \
      "$OURS_BIN" --ignored --nocapture >/dev/null

  gen_flamegraph "$OURS_LAND_PERF" "$OURS_LAND_SVG" \
    "stage2 ours landing ${TYPE}.${SAMPLE}" "iters=${ITERS} warmup=${WARMUP} freq=${FREQ}"
  gen_hotspots "$OURS_LAND_PERF" "$OURS_LAND_HOT"
fi

# Routinator baseline decode
ROUT_DECODE_PERF="$OUT_DIR/perf_stage2_routinator_${TYPE}_${SAMPLE}_decode.data"
ROUT_DECODE_SVG="$OUT_DIR/flamegraph_stage2_routinator_${TYPE}_${SAMPLE}_decode_release.svg"
ROUT_DECODE_HOT="$OUT_DIR/hotspots_stage2_routinator_${TYPE}_${SAMPLE}_decode_release.txt"

run_perf_profile "$ROUT_DECODE_PERF" \
  "$ROUT_BIN" \
    --dir "$ROOT_DIR/tests/benchmark/selected_der_v2" \
    --type "$TYPE" \
    --sample "$SAMPLE" \
    --iters "$ITERS" \
    --rounds 1 \
    --warmup-iters "$WARMUP" \
    --min-round-ms 1 \
    --max-iters 1000000 \
    --strict true \
    --out-csv "$OUT_DIR/_tmp_routinator_profile.csv" \
    --out-md "$OUT_DIR/_tmp_routinator_profile.md" \
    >/dev/null

gen_flamegraph "$ROUT_DECODE_PERF" "$ROUT_DECODE_SVG" \
  "stage2 routinator decode ${TYPE}.${SAMPLE}" "iters=${ITERS} warmup=${WARMUP} freq=${FREQ}"
gen_hotspots "$ROUT_DECODE_PERF" "$ROUT_DECODE_HOT"

echo "Done." >&2
echo "- ours decode flamegraph: $OURS_DECODE_SVG" >&2
echo "- ours decode hotspots:  $OURS_DECODE_HOT" >&2
if [[ "${PROFILE_LANDING}" == "1" ]]; then
  echo "- ours landing flamegraph: $OURS_LAND_SVG" >&2
  echo "- ours landing hotspots:  $OURS_LAND_HOT" >&2
fi
echo "- rout decode flamegraph: $ROUT_DECODE_SVG" >&2
echo "- rout decode hotspots:  $ROUT_DECODE_HOT" >&2

