#!/usr/bin/env bash

set -euo pipefail

BENCH=${BENCH:-build-ninja-arm64-metal-release/bin/llama-bench}
RESULTS_DIR=${RESULTS_DIR:-prefill-bench-results}
RUN_INDEX="${RESULTS_DIR}/runs.csv"
MODELS=(
    "LLMs/Llama-3.2-3B-Instruct-Q8_0.gguf"
    "LLMs/Phi-3.5-mini-instruct-Q8_0.gguf"
    "LLMs/Qwen3-4B-Q8_0.gguf"
)

REPETITIONS=100
PROMPT_VALUES=(512 2048)
BATCH_VALUES=(2048)
UBATCH_VALUES=(512 2048)


if [[ ! -x "$BENCH" ]]; then
    printf 'Could not find executable "%s".\n' "$BENCH" >&2
    exit 1
fi

mkdir -p "$RESULTS_DIR"

printf 'Running prefill benchmarks...\n'
printf 'Models: %s\n' "${MODELS[*]}"
printf 'Results directory: %s\n' "$RESULTS_DIR"
printf 'Prompt sweep: %s\n' "${PROMPT_VALUES[*]}"
printf 'Batch sweep: %s\n' "${BATCH_VALUES[*]}"
printf 'UBatch sweep: %s\n' "${UBATCH_VALUES[*]}"
printf 'Skipping combinations where ubatch is greater than batch.\n'
printf 'Output tokens: 0\n\n'

printf 'model,prompt_tokens,batch_size,ubatch_size,n_gen,repetitions,bench_json\n' > "$RUN_INDEX"

for MODEL in "${MODELS[@]}"; do
    MODEL_NAME=${MODEL##*/}
    MODEL_NAME=${MODEL_NAME%.*}
    for P in "${PROMPT_VALUES[@]}"; do
        for B in "${BATCH_VALUES[@]}"; do
            for U in "${UBATCH_VALUES[@]}"; do
                if (( U > B )); then
                    continue
                fi

                RUN_NAME="${MODEL_NAME}_p${P}_b${B}_ub${U}_n0_r${REPETITIONS}"
                BENCH_JSON="${RESULTS_DIR}/${RUN_NAME}.json"

                printf '"%s",%s,%s,%s,0,%s,"%s"\n' \
                    "${MODEL//\"/\"\"}" "$P" "$B" "$U" "$REPETITIONS" \
                    "${BENCH_JSON//\"/\"\"}" >> "$RUN_INDEX"
                printf 'Running %s pp %s with -b %s -ub %s...\n' "$MODEL_NAME" "$P" "$B" "$U"

                if "$BENCH" \
                    -m "$MODEL" \
                    -ngl -1 \
                    -sm none \
                    -mg 0 \
                    -r "$REPETITIONS" \
                    -p "$P" \
                    -n 0 \
                    -b "$B" \
                    -ub "$U" \
                    -o json \
                    --progress > "$BENCH_JSON"; then
                    :
                else
                    STATUS=$?
                    printf 'Benchmark failed with exit code %s.\n' "$STATUS" >&2
                    printf 'Last benchmark JSON: "%s"\n' "$BENCH_JSON" >&2
                    exit "$STATUS"
                fi
            done
        done
    done
done

printf 'All benchmarks complete.\n'
printf 'Run index written to "%s".\n' "$RUN_INDEX"
printf 'Per-run JSON files written to "%s".\n' "$RESULTS_DIR"