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

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
ROOT_DIR="$(cd -- "$SCRIPT_DIR/.." && pwd)"
BUILD_DIR="${SLIME_BENCH_BUILD_DIR:-/tmp/slime-bench-build}"
OUT_DIR="${SLIME_BENCH_OUT_DIR:-/tmp/slime-bench-results}"

CXX="${CXX:-g++}"
NVCC="${NVCC:-nvcc}"
CXXFLAGS="${CXXFLAGS:--O3 -std=c++17 -march=native -pthread -fwrapv}"
NVCCFLAGS="${NVCCFLAGS:--O3 -std=c++17 -Xcompiler=-fwrapv}"

RUNS=3
WARMUPS=1
DO_CPU=1
DO_CUDA=1
KEEP_OUTPUT=0
SEARCH_ARGS=()

usage() {
    cat <<'EOF'
Usage: bash codex/bench_wsl.sh [bench options] [-- common slime args]

Bench options:
  --runs <n>          Timed runs per binary (default: 3)
  --warmups <n>       Untimed warmup runs per binary (default: 1)
  --skip-cpu          Do not build/run CPU binaries
  --skip-cuda         Do not build/run CUDA binaries
  --build-dir <dir>   Build directory (default: /tmp/slime-bench-build)
  --out-dir <dir>     Output directory (default: /tmp/slime-bench-results)
  --keep-output       Keep per-run stdout/stderr logs
  --help              Show this help

Common slime args are passed to all binaries. Use only options supported by both
the original programs and the codex programs: -s, -e, -n, -d, -r.

Default common args:
  -s 0 -e 100000 -n 3 -r 60

Examples:
  bash codex/bench_wsl.sh
  bash codex/bench_wsl.sh --runs 5 -- -s 0 -e 1000000 -n 3 -r 100
  bash codex/bench_wsl.sh --skip-cuda -- -s 0 -e 200000 -n 4 -r 80

Notes:
  * CPU comparison runs codex/slime_cpu with --all-locations to match the
    original CPU program, which scans every matching location.
  * CUDA comparison runs codex/slime_cuda in first-hit-per-seed mode to match
    the original CUDA program.
  * The original CPU program ignores non-zero -s while splitting thread ranges,
    so use -s 0 for fair CPU result counts.
EOF
}

die() {
    echo "error: $*" >&2
    exit 1
}

need_value() {
    local name="$1"
    local value="${2:-}"
    [[ -n "$value" ]] || die "$name needs a value"
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --runs)
            need_value "$1" "${2:-}"
            RUNS="$2"
            shift 2
            ;;
        --warmups)
            need_value "$1" "${2:-}"
            WARMUPS="$2"
            shift 2
            ;;
        --skip-cpu)
            DO_CPU=0
            shift
            ;;
        --skip-cuda)
            DO_CUDA=0
            shift
            ;;
        --build-dir)
            need_value "$1" "${2:-}"
            BUILD_DIR="$2"
            shift 2
            ;;
        --out-dir)
            need_value "$1" "${2:-}"
            OUT_DIR="$2"
            shift 2
            ;;
        --keep-output)
            KEEP_OUTPUT=1
            shift
            ;;
        --help)
            usage
            exit 0
            ;;
        --)
            shift
            SEARCH_ARGS=("$@")
            break
            ;;
        *)
            SEARCH_ARGS+=("$1")
            shift
            ;;
    esac
done

[[ "$RUNS" =~ ^[0-9]+$ && "$RUNS" -gt 0 ]] || die "--runs must be a positive integer"
[[ "$WARMUPS" =~ ^[0-9]+$ ]] || die "--warmups must be a non-negative integer"

if [[ "${#SEARCH_ARGS[@]}" -eq 0 ]]; then
    SEARCH_ARGS=(-s 0 -e 10000000 -n 3 -r 60)
fi

START_SEED_ARG=0
for ((i = 0; i < ${#SEARCH_ARGS[@]}; ++i)); do
    if [[ "${SEARCH_ARGS[$i]}" == "-s" && $((i + 1)) -lt ${#SEARCH_ARGS[@]} ]]; then
        START_SEED_ARG="${SEARCH_ARGS[$((i + 1))]}"
    fi
done

mkdir -p "$BUILD_DIR" "$OUT_DIR"

read -r -a CXXFLAGS_ARR <<< "$CXXFLAGS"
read -r -a NVCCFLAGS_ARR <<< "$NVCCFLAGS"

build_cpu() {
    command -v "$CXX" >/dev/null 2>&1 || die "cannot find C++ compiler '$CXX'"

    echo "Building CPU binaries with: $CXX $CXXFLAGS"
    "$CXX" "${CXXFLAGS_ARR[@]}" "$ROOT_DIR/slime.cpp" -o "$BUILD_DIR/user_cpu"
    "$CXX" "${CXXFLAGS_ARR[@]}" "$ROOT_DIR/codex/slime_cpu.cpp" -o "$BUILD_DIR/codex_cpu"
}

build_cuda() {
    if ! command -v "$NVCC" >/dev/null 2>&1; then
        echo "Skipping CUDA: cannot find '$NVCC'" >&2
        DO_CUDA=0
        return
    fi

    echo "Building CUDA binaries with: $NVCC $NVCCFLAGS"
    "$NVCC" "${NVCCFLAGS_ARR[@]}" "$ROOT_DIR/slime.cu" -o "$BUILD_DIR/user_cuda"
    "$NVCC" "${NVCCFLAGS_ARR[@]}" "$ROOT_DIR/codex/slime_cuda.cu" -o "$BUILD_DIR/codex_cuda"
}

extract_count() {
    local file="$1"
    local count
    count="$(sed -n 's/.*hits=\([0-9][0-9]*\).*/\1/p' "$file" | tail -n 1)"
    if [[ -z "$count" ]]; then
        count="$(sed -n 's/.*found \([0-9][0-9]*\) seeds.*/\1/p' "$file" | tail -n 1)"
    fi
    if [[ -z "$count" ]]; then
        count="$(grep -Ec '^(hit|found seed)' "$file" || true)"
    fi
    echo "${count:-unknown}"
}

run_once() {
    local name="$1"
    shift

    local stdout_file="$OUT_DIR/$name.stdout"
    local stderr_file="$OUT_DIR/$name.stderr"
    local time_file="$OUT_DIR/$name.time"

    set +e
    /usr/bin/time -f '%e' -o "$time_file" "$@" >"$stdout_file" 2>"$stderr_file"
    local status=$?
    set -e

    if [[ "$status" -ne 0 ]]; then
        echo "Command failed: $*" >&2
        echo "--- stderr ---" >&2
        tail -n 80 "$stderr_file" >&2 || true
        exit "$status"
    fi

    local elapsed
    elapsed="$(tr -d '\r\n' < "$time_file")"
    local count
    count="$(extract_count "$stdout_file")"
    echo "$elapsed $count"
}

mean() {
    awk '
        NF { sum += $1; n += 1 }
        END {
            if (n == 0) print "nan";
            else printf "%.6f", sum / n;
        }
    '
}

speedup() {
    local base="$1"
    local candidate="$2"
    awk -v base="$base" -v candidate="$candidate" '
        BEGIN {
            if (candidate > 0) printf "%.2fx", base / candidate;
            else printf "n/a";
        }
    '
}

bench_binary() {
    local label="$1"
    shift
    local -a command=("$@")

    echo
    echo "== $label =="
    echo "Command: ${command[*]}"

    for ((i = 1; i <= WARMUPS; ++i)); do
        echo "Warmup $i/$WARMUPS..."
        run_once "${label// /_}.warmup.$i" "${command[@]}" >/dev/null
    done

    local times_file="$OUT_DIR/${label// /_}.times"
    : > "$times_file"
    local last_count="unknown"
    for ((i = 1; i <= RUNS; ++i)); do
        local result elapsed count
        result="$(run_once "${label// /_}.run.$i" "${command[@]}")"
        elapsed="${result%% *}"
        count="${result#* }"
        last_count="$count"
        echo "$elapsed" >> "$times_file"
        printf 'Run %d/%d: %ss, reported=%s\n' "$i" "$RUNS" "$elapsed" "$count"
    done

    local avg
    avg="$(mean < "$times_file")"
    printf '%s\t%s\t%s\n' "$label" "$avg" "$last_count" >> "$OUT_DIR/summary.tsv"
}

print_summary() {
    echo
    echo "Summary"
    printf '%-16s %12s %12s %12s\n' "binary" "avg_s" "reported" "speedup"

    local user_cpu_avg="" codex_cpu_avg="" user_cuda_avg="" codex_cuda_avg=""
    while IFS=$'\t' read -r label avg count; do
        case "$label" in
            user_cpu) user_cpu_avg="$avg" ;;
            codex_cpu) codex_cpu_avg="$avg" ;;
            user_cuda) user_cuda_avg="$avg" ;;
            codex_cuda) codex_cuda_avg="$avg" ;;
        esac
    done < "$OUT_DIR/summary.tsv"

    while IFS=$'\t' read -r label avg count; do
        local base="" sp="baseline"
        case "$label" in
            codex_cpu) base="$user_cpu_avg" ;;
            codex_cuda) base="$user_cuda_avg" ;;
        esac
        if [[ -n "$base" ]]; then
            sp="$(speedup "$base" "$avg")"
        fi
        printf '%-16s %12s %12s %12s\n' "$label" "$avg" "$count" "$sp"
    done < "$OUT_DIR/summary.tsv"
}

echo "Root:       $ROOT_DIR"
echo "Build dir:  $BUILD_DIR"
echo "Output dir: $OUT_DIR"
echo "Args:       ${SEARCH_ARGS[*]}"
echo "Runs:       $RUNS timed, $WARMUPS warmup"

if [[ "$DO_CPU" -eq 1 && "$START_SEED_ARG" != "0" ]]; then
    echo "Warning: original slime.cpp ignores non-zero -s in its CPU thread split; CPU counts will not be apples-to-apples." >&2
fi

: > "$OUT_DIR/summary.tsv"

if [[ "$DO_CPU" -eq 1 ]]; then
    build_cpu
    bench_binary user_cpu "$BUILD_DIR/user_cpu" "${SEARCH_ARGS[@]}"
    bench_binary codex_cpu "$BUILD_DIR/codex_cpu" "${SEARCH_ARGS[@]}" --all-locations
fi

if [[ "$DO_CUDA" -eq 1 ]]; then
    build_cuda
    if [[ "$DO_CUDA" -eq 1 ]]; then
        bench_binary user_cuda "$BUILD_DIR/user_cuda" "${SEARCH_ARGS[@]}"
        bench_binary codex_cuda "$BUILD_DIR/codex_cuda" "${SEARCH_ARGS[@]}"
    fi
fi

print_summary

if [[ "$KEEP_OUTPUT" -eq 0 ]]; then
    find "$OUT_DIR" -type f \( -name '*.stdout' -o -name '*.stderr' -o -name '*.time' -o -name '*.times' \) -delete
else
    echo
    echo "Kept logs in $OUT_DIR"
fi