#!/bin/bash

MODEL_PATH=$1
PARSER=$2
PORT=${3:-8000}

# find the last folder name of the model path
MODEL_NAME=$(basename $MODEL_PATH)

# Start vllm serve in the background
uv run --with vllm vllm serve $MODEL_PATH --uvicorn-log-level error --disable-log-stats --disable-log-requests --disable-fastapi-docs --reasoning-parser deepseek_r1 --tool-call-parser $PARSER --enable-auto-tool-choice --port $PORT &

# Capture the PID of the vllm process
VLLM_PID=$!
echo "vllm serve process started with PID $VLLM_PID"

# Function to check the health endpoint
check_health() {
  while true; do
    # Use `curl` to send a request to the health endpoint
    if response=$(curl -s -o /dev/null -w "%{http_code}" http://localhost:$PORT/health); then
      if [ "$response" -eq 200 ]; then
        echo "Health check passed with status 200."
        return 0
      else
        echo "Health check failed with status $response. Retrying in 10 seconds..."
      fi
    else
      echo "Failed to connect to the health endpoint. Retrying in 10 seconds..."
    fi
    sleep 10
  done
}

(check_health)

# Runs benchmark to runs/model_name/local
# Create an array to store PIDs
declare -a PIDS

Run each command in the background and store its PID
LOCAL_LLM_PORT=$PORT uv run python -m agentdojo.scripts.benchmark --model local --attack important_instructions --logdir "./runs-final/${MODEL_NAME}" --suite workspace &
PIDS+=($!)

LOCAL_LLM_PORT=$PORT uv run python -m agentdojo.scripts.benchmark --model local --attack important_instructions --logdir "./runs-final/${MODEL_NAME}" --suite slack &
PIDS+=($!)

LOCAL_LLM_PORT=$PORT uv run python -m agentdojo.scripts.benchmark --model local --attack important_instructions --logdir "./runs-final/${MODEL_NAME}" --suite travel &
PIDS+=($!)

LOCAL_LLM_PORT=$PORT uv run python -m agentdojo.scripts.benchmark --model local --attack important_instructions --logdir "./runs-final/${MODEL_NAME}" --suite banking &
PIDS+=($!)

echo "Started processes with PIDs: ${PIDS[@]}"

# Wait for each specific PID
for pid in "${PIDS[@]}"; do
    wait "$pid"
    echo "Process $pid finished."
done

# After the script completes, terminate the vllm process
if ps -p $VLLM_PID > /dev/null; then
  echo "Stopping vllm serve process..."
  kill $VLLM_PID
  wait $VLLM_PID 2>/dev/null
  echo "vllm serve process stopped."
  sleep 60
else
  echo "vllm serve process is not running."
fi
