#!/bin/sh
# Reproducibly construct an exact-token canonical Gemma user prompt.
set -eu

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

model=$1
target=$2
output=$3
case $target in
    ''|*[!0-9]*|0) echo "gem: target token count must be a positive integer" >&2; exit 2 ;;
esac
if [ ! -r "$model" ]; then
    echo "gem: model must be readable: $model" >&2
    exit 2
fi
if [ -e "$output" ]; then
    echo "gem: output already exists: $output" >&2
    exit 2
fi
if ! command -v jq >/dev/null 2>&1 || [ ! -x build/gem-tokenize ]; then
    echo "gem: jq and build/gem-tokenize are required" >&2
    exit 2
fi

probe_one=$(mktemp /private/tmp/gem-long-prompt-one.XXXXXX)
probe_two=$(mktemp /private/tmp/gem-long-prompt-two.XXXXXX)
candidate=$(mktemp /private/tmp/gem-long-prompt-candidate.XXXXXX)
cleanup()
{
    rm -f "$probe_one" "$probe_two" "$candidate"
}
trap cleanup EXIT INT TERM HUP

base_line='context datum 314159'
awk -v line="$base_line" 'BEGIN { print line }' > "$probe_one"
awk -v line="$base_line" 'BEGIN { print line; print line }' > "$probe_two"
one_count=$(build/gem-tokenize --count --file --chat-user "$model" "$probe_one" |
    jq -er '.token_count')
two_count=$(build/gem-tokenize --count --file --chat-user "$model" "$probe_two" |
    jq -er '.token_count')
line_tokens=$((two_count - one_count))
fixed_tokens=$((one_count - line_tokens))
if [ "$line_tokens" -le 0 ] || [ "$fixed_tokens" -lt 0 ] ||
   [ "$target" -lt "$fixed_tokens" ]; then
    echo "gem: tokenizer did not produce a usable repeated-line prompt shape" >&2
    exit 1
fi

full_lines=$(((target - fixed_tokens) / line_tokens))
remainder=$((target - fixed_tokens - full_lines * line_tokens))
awk -v line="$base_line" -v lines="$full_lines" -v remainder="$remainder" '
BEGIN {
    for (i = 0; i < lines; ++i) print line
    if (remainder == 1) {
        print ""
    } else if (remainder > 1) {
        split("x y z q r s t u v", words, " ")
        for (i = 1; i < remainder; ++i)
            printf "%s%s", (i == 1 ? "" : " "), words[i]
        printf "\n"
    }
}' > "$candidate"

actual=$(build/gem-tokenize --count --file --chat-user "$model" "$candidate" |
    jq -er '.token_count')
if [ "$actual" -ne "$target" ]; then
    echo "gem: generated prompt has $actual chat tokens; requested $target" >&2
    exit 1
fi
mv "$candidate" "$output"
printf '{"schema_version":1,"chat_tokens":%s,"bytes":%s,"sha256":"%s"}\n' \
    "$actual" "$(wc -c < "$output" | tr -d ' ')" \
    "$(shasum -a 256 "$output" | awk '{print $1}')"