#!/bin/sh

# Usage: ./test_prompts.sh [-i prompt_file] [-o result_file]

SOCK_PATH="/var/run/lexd/sock"

# Defaults
PROMPT_FILE="test_prompts.txt"
RESULT_FILE="results.txt"

# Parse -i and -o flags
while getopts "i:o:" opt; do
    case "$opt" in
        i) PROMPT_FILE="$OPTARG" ;;
        o) RESULT_FILE="$OPTARG" ;;
        *) echo "Usage: $0 [-i prompt_file] [-o result_file]" >&2; exit 1 ;;
    esac
done

# Ensure socket exists (daemon must already be running)
if [ ! -S "$SOCK_PATH" ]; then
    echo "Error: Socket '$SOCK_PATH' not found. Is lexd running?" >&2
    exit 1
fi

# Retrieve running daemon PID
DAEMON_PID=$(pgrep -f "lexd" | head -n 1)

if [ -z "$DAEMON_PID" ]; then
    echo "Warning: Could not detect DAEMON_PID. RSS metrics may be unavailable." >&2
fi

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

# Helper function to get current swap occupancy in MB by parsing swapctl -l
get_swap_occupancy_mb() {
    swapctl -l | awk '
        NR > 1 {
            used_blocks += $3
        }
        END {
            printf "%.2f", used_blocks * 512 / 1048576
        }'
}

# Count valid prompts using POSIX awk
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) result file
> "$RESULT_FILE"

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

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

cleanup() {
    rm -f "$TIME_LOG"
}
trap cleanup EXIT INT TERM

echo "Evaluating $total_prompts prompt(s) from '$PROMPT_FILE'"
echo "Writing results to: '$RESULT_FILE'"
echo "---------------------------------------------------"

current=0
first=1

# Process each prompt
while IFS= read -r prompt || [ -n "$prompt" ]; do
    case "$prompt" in
        ""|\#*) continue ;;
    esac

    current=$((current + 1))

    if [ "$first" -eq 1 ]; then
        first=0
    else
        printf "%s\n\n" "$RULE" >> "$RESULT_FILE"
    fi

    printf "PROMPT: %s\n\n" "$prompt" >> "$RESULT_FILE"

    # Sample daemon RSS prior to request
    if [ -n "$DAEMON_PID" ]; then
        daemon_rss_kb=$(ps -o rss= -p "$DAEMON_PID" 2>/dev/null | tr -d ' ')
    fi

    # Capture baseline swap occupancy before sending the prompt
    swap_start_mb=$(get_swap_occupancy_mb)

    # Send prompt to daemon over socket (exact original invocation)
    /usr/bin/time -lp sh -c "printf '%s\n' \"\$1\" | nc -U \"\$2\"" _ "$prompt" "$SOCK_PATH" \
        >> "$RESULT_FILE" 2> "$TIME_LOG"

    # Capture ending swap occupancy after prompt completes
    swap_end_mb=$(get_swap_occupancy_mb)

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

    # Calculate net change in swap occupancy during request execution
    if [ -n "$swap_start_mb" ] && [ -n "$swap_end_mb" ]; then
        swap_delta_mb=$(awk -v start="$swap_start_mb" -v end="$swap_end_mb" \
            'BEGIN { printf "%+.2f", end - start }')
    else
        swap_end_mb="unavailable"
        swap_delta_mb="+0.00"
    fi

    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 | SWAP: ${swap_end_mb} MB (${swap_delta_mb} MB)"
    else
        metrics_summary="ELAPSED: ${real_time}s | DAEMON RSS: unavailable | SWAP OCCUPIED: ${swap_end_mb} MB (${swap_delta_mb} MB)"
    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)."