#!/bin/sh
set -eu

if [ "$#" -ne 2 ]; then
    echo "usage: $0 GENERATOR MODEL.gguf" >&2
    exit 2
fi

generator=$1
model=$2
state_dir=$(mktemp -d "${TMPDIR:-/tmp}/gem-long-prompt-test.XXXXXX")
cleanup()
{
    rm -f "$state_dir"/prompt-*.txt "$state_dir"/receipt-*.json
    rmdir "$state_dir"
}
trap cleanup EXIT INT TERM HUP

for target in 8192 32768 65536 131072; do
    prompt=$state_dir/prompt-$target.txt
    receipt=$state_dir/receipt-$target.json
    "$generator" "$model" "$target" "$prompt" > "$receipt"
    jq -e --argjson target "$target" '
        .schema_version == 1 and .chat_tokens == $target and
        (.bytes | type) == "number" and .bytes > 0 and
        (.sha256 | test("^[0-9a-f]{64}$"))
    ' "$receipt" >/dev/null
    counted=$(build/gem-tokenize --count --file --chat-user "$model" "$prompt" |
        jq -er '.token_count')
    test "$counted" -eq "$target"
done

test "$(jq -r '.sha256' "$state_dir/receipt-8192.json")" = \
    e652edc2bea3a7fbd58a788e024cb23d97f0be1b734c7575963da130726c74dd
test "$(jq -r '.sha256' "$state_dir/receipt-32768.json")" = \
    21a5c34222164c5261e47f5c3d2f9e0ff17cb7fb110705661445c404b99073b8
test "$(jq -r '.sha256' "$state_dir/receipt-65536.json")" = \
    b7fe45961d8232f18b35b53567d132936812805a486bf59acb0332ed5f0a7e3d
test "$(jq -r '.sha256' "$state_dir/receipt-131072.json")" = \
    a0f001d27776bc653c1c44efd71728f042b6e3686a7705ae23d00f9900f70c32