#!/bin/sh

# Usage: ./test_prompts.sh <model_path> [prompt_file]

SOCK_PATH="/var/run/lex.sock"
EXECUTABLE="./lex"

# Mandatory Model Argument
MODEL_PATH="$1"

if [ -z "$MODEL_PATH" ]; then
    echo "Usage: $0 <model_path> [prompt_file]" >&2
    exit 1
fi

if [ ! -f "$MODEL_PATH" ]; then
    echo "Error: Model file '$MODEL_PATH' not found." >&2
    exit 1
fi

# Derive base model name and output result filename
# e.g., models/qwen2.5-7b-q4.gguf -> results_qwen2.5-7b-q4.gguf.txt
MODEL_NAME=$(basename "$MODEL_PATH")
RESULT_FILE="results_${MODEL_NAME}.txt"

# Default prompt filename (optional 2nd argument)
PROMPT_FILE="${2:-test_prompts.txt}"

# Ensure the executable exists and is runnable
if [ ! -x "$EXECUTABLE" ]; then
    echo "Error: Executable '$EXECUTABLE' not found or not executable." >&2
    exit 1
fi

# Ensure the prompt file exists
if [ ! -f "$PROMPT_FILE" ]; then
    echo "Error: Prompt file '$PROMPT_FILE' not found." >&2
    exit 1
fi

# Count valid prompts using POSIX awk to avoid ksh $(...) parenthetical syntax errors
total_prompts=$(awk '/^[[:space:]]*($|#)/ {next} {count++} END {print count+0}' "$PROMPT_FILE")

if [ "$total_prompts" -eq 0 ]; then
    echo "No valid prompts found in '$PROMPT_FILE'."
    exit 0
fi

# Create (or truncate) the result file at the start of the run
> "$RESULT_FILE"

# Create horizontal rule.
RULE=""
i=0
while [ $i -lt 72 ]; do
    RULE="${RULE}-"
    i=$((i + 1))
done

# --- Start the daemon once, wait for it to come up, capture its baseline RSS ---

echo "Starting daemon with model: '$MODEL_PATH'"
"$EXECUTABLE" -m "$MODEL_PATH" &
DAEMON_PID=$!

cleanup() {
    kill "$DAEMON_PID" 2>/dev/null
    wait "$DAEMON_PID" 2>/dev/null
    rm -f "$TIME_LOG"
}
trap cleanup EXIT INT TERM

# Wait for the socket to appear (covers model load time) and for the
# process to still be alive (covers early startup failure, e.g. bad
# model path -> daemon exits immediately).
waited=0
while [ ! -S "$SOCK_PATH" ]; do
    if ! kill -0 "$DAEMON_PID" 2>/dev/null; then
        echo "Error: daemon exited during startup (bad model path?)" >&2
        exit 1
    fi
    sleep 1
    waited=$((waited + 1))
    if [ "$waited" -ge 120 ]; then
        echo "Error: timed out waiting for daemon socket at $SOCK_PATH" >&2
        exit 1
    fi
done

echo "Daemon ready (pid $DAEMON_PID) after ${waited}s"
echo "Evaluating $total_prompts prompt(s) from '$PROMPT_FILE'"
echo "Writing results to: '$RESULT_FILE'"
echo "---------------------------------------------------"

current=0
first=1

# Temporary file to capture time / rusage output
TIME_LOG=$(mktemp /tmp/lex_time.XXXXXX)

# Process each prompt line by line
while IFS= read -r prompt || [ -n "$prompt" ]; do
    # Skip empty lines or comment lines starting with #
    case "$prompt" in
        ""|\#*) continue ;;
    esac

    current=$((current + 1))

    # Add delimiter before every prompt except the first in the results file
    if [ "$first" -eq 1 ]; then
        first=0
    else
        printf "%s\n\n" "$RULE" >> "$RESULT_FILE"
    fi

    # Log prompt header to the result file
    printf "PROMPT: %s\n\n" "$prompt" >> "$RESULT_FILE"

    # Sample the daemon's RSS right before firing the request, so we have
    # a rough steady-state figure (won't catch a mid-request peak, but the
    # model's footprint is dominated by the loaded weights, not per-request
    # allocation, so this is a reasonable proxy).
    daemon_rss_kb=$(ps -o rss= -p "$DAEMON_PID" 2>/dev/null | tr -d ' ')

    # Send the prompt to the daemon over the socket, timing the round trip.
    # This now measures request latency only (no model-load cost), which is
    # the number that actually matters day to day with the daemon design.
    /usr/bin/time -lp sh -c "printf '%s\n' \"\$1\" | nc -U \"\$2\"" _ "$prompt" "$SOCK_PATH" \
        >> "$RESULT_FILE" 2> "$TIME_LOG"

    real_time=$(awk '/^real/ {print $2}' "$TIME_LOG")

    if [ -n "$daemon_rss_kb" ]; then
        daemon_rss_mb=$(awk -v kb="$daemon_rss_kb" 'BEGIN{printf "%.2f", kb/1024}')
        metrics_summary="ELAPSED: ${real_time}s | DAEMON RSS: ${daemon_rss_mb} MB"
    else
        metrics_summary="ELAPSED: ${real_time}s | DAEMON RSS: unavailable"
    fi

    printf "METRICS: %s\n\n" "$metrics_summary" >> "$RESULT_FILE"

    printf "[%d/%d] Done: %s (%s)\n" "$current" "$total_prompts" "$prompt" "$metrics_summary"

done < "$PROMPT_FILE"

echo "---------------------------------------------------"
echo "Done! Processed $total_prompts prompt(s)."