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

# Record baseline timings for C and Python (and Hakorune VM/AOT) into benchmarks/baselines/.
# Usage: record_baselines.sh <bench_key|all> [warmup=2] [repeat=7]
# Env:
#   PERF_SUBTRACT_STARTUP=1  subtract minimal startup baseline (ret0) for VM/AOT
#   NYASH_LLVM_BACKEND=crate|native  select LLVM builder backend for AOT (default auto)
#   bench_key: box_create_destroy_small | method_call_only_small | all

ROOT_DIR=$(cd "$(dirname "$0")/../.." && pwd)
TARGET_DIR="${ROOT_DIR}/target"
OUT_DIR="${ROOT_DIR}/benchmarks/baselines"
PY_DIR="${ROOT_DIR}/benchmarks/python"
C_DIR="${ROOT_DIR}/benchmarks/c"

KEY=${1:-all}
WARMUP=${2:-2}
REPEAT=${3:-7}

mkdir -p "${OUT_DIR}"

time_ms() { date +%s%3N; }
measure_cmd_ms() { local t1 t2; t1=$(time_ms); "$@" >/dev/null 2>&1 || true; t2=$(time_ms); echo $((t2-t1)); }
median_ms() { awk 'NF{print $1}' | sort -n | awk '{a[NR]=$1} END{ if(NR==0){print 0; exit} n=int((NR+1)/2); print a[n] }'; }

collect_series() {
  local warmup=$1; shift
  local repeat=$1; shift
  local -a cmd=("$@")
  for _ in $(seq 1 "${warmup}"); do measure_cmd_ms "${cmd[@]}" >/dev/null || true; done
  for _ in $(seq 1 "${repeat}"); do measure_cmd_ms "${cmd[@]}"; done
}

ensure_c_built() {
  local key=$1
  local c_src="${C_DIR}/bench_${key}.c"
  local c_bin="${TARGET_DIR}/perf_c_${key}"
  if [[ ! -f "${c_src}" ]]; then echo "[error] missing ${c_src}" >&2; return 1; fi
  mkdir -p "${TARGET_DIR}"
  cc -O3 -march=native -mtune=native -o "${c_bin}" "${c_src}" 2>/dev/null || cc -O3 -o "${c_bin}" "${c_src}"
}

record_one() {
  local key=$1
  local ts host c_ms py_ms ny_vm_ms ny_aot_ms

  # C
  ensure_c_built "${key}"
  local c_bin="${TARGET_DIR}/perf_c_${key}"
  local c_series; c_series=$(collect_series "${WARMUP}" "${REPEAT}" "${c_bin}")
  c_ms=$(printf "%s\n" "${c_series}" | median_ms)

  # Python
  local py_file="${PY_DIR}/bench_${key}.py"
  if command -v python3 >/dev/null 2>&1 && [[ -f "${py_file}" ]]; then
    local py_series; py_series=$(collect_series "${WARMUP}" "${REPEAT}" python3 "${py_file}")
    py_ms=$(printf "%s\n" "${py_series}" | median_ms)
  else
    py_ms=0
  fi

  # Hakorune VM (optional)
  local hako_bin="${TARGET_DIR}/release/hakorune"
  local hako_prog="${ROOT_DIR}/benchmarks/bench_${key}.hako"
  if [[ -x "${hako_bin}" && -f "${hako_prog}" ]]; then
    local -a HAKO_ENV=(NYASH_PARSER_STAGE3=1 HAKO_PARSER_STAGE3=1 NYASH_PARSER_ALLOW_SEMICOLON=1)
    local hako_series; hako_series=$(collect_series "${WARMUP}" "${REPEAT}" env "${HAKO_ENV[@]}" timeout 20s "${hako_bin}" --backend vm "${hako_prog}")
    ny_vm_ms=$(printf "%s\n" "${hako_series}" | median_ms)
    # Optional subtract: measure minimal VM startup (ret0)
    if [[ "${PERF_SUBTRACT_STARTUP:-0}" == "1" ]]; then
      local tmp_ret0="$(mktemp --suffix .hako)"; cat >"${tmp_ret0}" <<'HAKO'
static box Main { main() { return 0 } }
HAKO
      local base_series; base_series=$(collect_series 1 3 env "${HAKO_ENV[@]}" timeout 20s "${hako_bin}" --backend vm "${tmp_ret0}")
      local base_ms; base_ms=$(printf "%s\n" "${base_series}" | median_ms)
      rm -f "${tmp_ret0}" || true
      if [[ "${base_ms}" =~ ^[0-9]+$ ]]; then
        ny_vm_ms=$(( ny_vm_ms>base_ms ? ny_vm_ms-base_ms : 0 ))
      fi
    fi
  else
    ny_vm_ms=0
  fi

  # AOT (crate/native backend)
  ny_aot_ms=0
  if [[ -x "${hako_bin}" && -f "${hako_prog}" ]]; then
    # 1) Emit MIR JSON (prefer robust Rust CLI; fallback to Stage‑B wrapper)
    local tmp_json; tmp_json=$(mktemp --suffix .json)
    if "${hako_bin}" --emit-mir-json "${tmp_json}" "${hako_prog}" >/dev/null 2>&1 || \
       bash "${ROOT_DIR}/tools/hakorune_emit_mir.sh" "${hako_prog}" "${tmp_json}" >/dev/null 2>&1; then
      # 2) Build EXE via ny_mir_builder (default backend auto/crate/native)
      local exe_path="${TARGET_DIR}/perf_ny_${key}.exe"
      if bash "${ROOT_DIR}/tools/ny_mir_builder.sh" --in "${tmp_json}" --emit exe -o "${exe_path}" --quiet >/dev/null 2>&1; then
        # 3) Measure run time of EXE
        local exe_series; exe_series=$(collect_series "${WARMUP}" "${REPEAT}" timeout 20s "${exe_path}")
        ny_aot_ms=$(printf "%s\n" "${exe_series}" | median_ms)
        # Optional subtract: minimal AOT startup (ret0)
        if [[ "${PERF_SUBTRACT_STARTUP:-0}" == "1" ]]; then
          # Build ret0 EXE
          local tmp_ret0_hako tmp_ret0_json ret0_exe
          tmp_ret0_hako=$(mktemp --suffix .hako); cat >"${tmp_ret0_hako}" <<'HAKO'
static box Main { main() { return 0 } }
HAKO
          tmp_ret0_json=$(mktemp --suffix .json)
          ret0_exe="${TARGET_DIR}/perf_ny_ret0.exe"
          if bash "${ROOT_DIR}/tools/hakorune_emit_mir.sh" "${tmp_ret0_hako}" "${tmp_ret0_json}" >/dev/null 2>&1 \
             && bash "${ROOT_DIR}/tools/ny_mir_builder.sh" --in "${tmp_ret0_json}" --emit exe -o "${ret0_exe}" --quiet >/dev/null 2>&1; then
            local base_series; base_series=$(collect_series 1 3 timeout 20s "${ret0_exe}")
            local base_ms; base_ms=$(printf "%s\n" "${base_series}" | median_ms)
            if [[ "${base_ms}" =~ ^[0-9]+$ ]]; then
              ny_aot_ms=$(( ny_aot_ms>base_ms ? ny_aot_ms-base_ms : 0 ))
            fi
          fi
          rm -f "${tmp_ret0_hako}" "${tmp_ret0_json}" || true
        fi
      fi
    fi
    rm -f "${tmp_json}" || true
  fi

  ts=$(date -Is)
  host=$(hostname 2>/dev/null || echo unknown)

  local obj; obj=$(cat <<JSON
{
  "bench": "${key}",
  "ts": "${ts}",
  "host": "${host}",
  "unit": "ms",
  "warmup": ${WARMUP},
  "repeat": ${REPEAT},
  "c_ms": ${c_ms},
  "py_ms": ${py_ms},
  "ny_vm_ms": ${ny_vm_ms},
  "ny_aot_ms": ${ny_aot_ms}
}
JSON
)

  printf "%s\n" "${obj}" > "${OUT_DIR}/${key}.latest.json"
  printf "%s\n" "${obj}" >> "${OUT_DIR}/${key}.ndjson"
  echo "[saved] ${OUT_DIR}/${key}.latest.json"
}

run_keys=("${KEY}")
if [[ "${KEY}" == "all" ]]; then
  run_keys=(box_create_destroy_small method_call_only_small)
fi

for k in "${run_keys[@]}"; do
  record_one "${k}"
done