#!/bin/bash

# Usage: ./ollama_query_runner.sh <avg_queries_per_minute> <num_queries>
# Example: ./ollama_query_runner.sh 10 0

AVG_QPM=$1
NUM_QUERIES=$2

if [[ -z "$AVG_QPM" ]]; then
  echo "Usage: $0 <avg_queries_per_minute> <num_queries (0 = infinite)>"
  exit 1
fi

#Install the dependencies
export DEBIAN_FRONTEND=noninteractive
sudo curl -fsSL https://ollama.com/install.sh | sh
nohup ollama serve > ollama.log 2>&1 &
sudo apt install -y bc

# List of countries
COUNTRIES=("France" "Germany" "Brazil" "Japan" "India" "Canada" "Egypt" "Australia" "Kenya" "Norway")

get_random_delay() {
  MIN_DELAY=$(echo "scale=2; (60 / $AVG_QPM) * 0.5" | bc)
  MAX_DELAY=$(echo "scale=2; (60 / $AVG_QPM) * 1.5" | bc)
  awk -v min="$MIN_DELAY" -v max="$MAX_DELAY" 'BEGIN { srand(); print min + rand() * (max - min) }'
}

get_random_country() {
  index=$((RANDOM % ${#COUNTRIES[@]}))
  echo "${COUNTRIES[$index]}"
}

COUNTER=0

while [[ "$NUM_QUERIES" -eq 0 || "$COUNTER" -lt "$NUM_QUERIES" ]]; do
  COUNTRY=$(get_random_country)
  QUERY="What is the capital of $COUNTRY?"
  echo "Submitting query: $QUERY"
  echo "$QUERY" | ollama run llama3
  COUNTER=$((COUNTER + 1))
  SLEEP_TIME=$(get_random_delay)
  echo "Sleeping for $SLEEP_TIME seconds..."
  sleep "$SLEEP_TIME"
done

