#!/usr/bin/env bash

set -euo pipefail

ROOT_DIR=$(cd "/Users/harmankhunguradevagent/C++projects/llama.cpp" && pwd)
MODEL_DIR=${MODEL_DIR:-"${ROOT_DIR}/LLMs"}
BENCH_BIN=${BENCH_BIN:-"${ROOT_DIR}/build-ninja-arm64-metal-release/bin/llama-bench"}
OUTPUT_FILE=${OUTPUT_FILE:-"${ROOT_DIR}/llama-bench-results.json"}

BENCH_ARGS=(
    --repetitions 100
    --n-prompt 512
    --n-gen 0,3277,6554,9830,13107
    --batch-size 2048
    --ubatch-size 512
    --n-gpu-layers -1
    --flash-attn auto
    --progress
)

if [[ ! -x "${BENCH_BIN}" ]]; then
    printf 'llama-bench executable not found: %s\n' "${BENCH_BIN}" >&2
    printf 'Set BENCH_BIN to the path of your llama-bench executable.\n' >&2
    exit 1
fi

if [[ ! -d "${MODEL_DIR}" ]]; then
    printf 'Model directory not found: %s\n' "${MODEL_DIR}" >&2
    exit 1
fi

families=(llama-3.2-3b phi-3.5-mini qwen3-4b)
family_patterns=('Llama-3.2-3B-*.gguf' 'Phi-3.5-*.gguf' 'Qwen3-4B-*.gguf')

output_dir=$(dirname "${OUTPUT_FILE}")

found_models=0
for family_index in "${!families[@]}"; do
    family=${families[${family_index}]}
    family_pattern=${family_patterns[${family_index}]}
    model_args=()
    while IFS= read -r -d '' model; do
        model_args+=(--model "${model}")
    done < <(find "${MODEL_DIR}" -type f -name "${family_pattern}" -print0)

    if (( ${#model_args[@]} == 0 )); then
        printf 'No models found for family %s; skipping\n' "${family}" >&2
        continue
    fi
    found_models=1

    output_file="${output_dir}/${family}-bench-results.json"
    output_tmp=$(mktemp "${output_file}.tmp.XXXXXX")
    trap 'rm -f "${output_tmp}"' EXIT

    printf 'Benchmarking family %s (%d model(s)). Results will be written to %s\n' \
        "${family}" "$(( ${#model_args[@]} / 2 ))" "${output_file}" >&2

    "${BENCH_BIN}" "${model_args[@]}" "${BENCH_ARGS[@]}" "$@" --output json > "${output_tmp}"
    mv "${output_tmp}" "${output_file}"
    trap - EXIT

    printf 'Benchmark complete: %s\n' "${output_file}" >&2
done

if (( found_models == 0 )); then
    printf 'No supported .gguf models found under: %s\n' "${MODEL_DIR}" >&2
    exit 1
fi