#!/bin/bash
# Staircase load test against the tokamak public API (prod or dev).
# Usage: ./scripts/tokamak-load.sh <start-rps> <step-rps> <step-duration> <steps> [env]
# Example (gentle first run):
#   ./scripts/tokamak-load.sh 2 2 1m 4
set -euo pipefail

START=${1:?start rps required}
STEP=${2:?step rps required}
DUR=${3:?step duration required, e.g. 2m}
STEPS=${4:?number of steps required}
ENVIRONMENT=${5:-prod}
NAME="tokamak-$(date +%s)"
PROM_RW="http://prometheus.monitoring.svc:9090/api/v1/write"

case "$ENVIRONMENT" in
  prod) BASE="https://api.blueocean.ai/public/v1" ;;
  dev)  BASE="https://api.dev.blueocean.ai/public/v1" ;;
  *) echo "env must be prod or dev"; exit 1 ;;
esac

kubectl -n loadtest apply -f - <<EOF
apiVersion: batch/v1
kind: Job
metadata:
  name: $NAME
spec:
  parallelism: 1
  ttlSecondsAfterFinished: 1800
  template:
    spec:
      restartPolicy: Never
      containers:
        - name: k6
          image: grafana/k6:1.4.0
          command: ["k6", "run", "/scripts/tokamak.js", "-o", "experimental-prometheus-rw", "--tag", "testid=$NAME"]
          env:
            - {name: BASE_URL, value: "$BASE"}
            - {name: START_RPS, value: "$START"}
            - {name: STEP_RPS, value: "$STEP"}
            - {name: STEP_DURATION, value: "$DUR"}
            - {name: STEPS, value: "$STEPS"}
            - {name: K6_PROMETHEUS_RW_SERVER_URL, value: "$PROM_RW"}
            - {name: K6_PROMETHEUS_RW_TREND_STATS, value: "avg,p(50),p(95),p(99),min,max"}
            - name: TOKAMAK_API_KEY
              valueFrom: {secretKeyRef: {name: tokamak-api-key, key: api-key}}
          resources:
            requests: {cpu: 200m, memory: 256Mi}
            limits: {memory: 1Gi}
          volumeMounts:
            - {name: scripts, mountPath: /scripts}
      volumes:
        - name: scripts
          configMap: {name: k6-tokamak-script}
EOF

MAX=$(( START + (STEPS - 1) * STEP ))
echo "tokamak staircase '$NAME' against $ENVIRONMENT ($BASE)"
echo "  $STEPS steps: $START -> $MAX rps (+$STEP each), $DUR holds"
echo "  auth: k8s secret tokamak-api-key"
echo "Watch:  Grafana -> 'k6 Step Load' dashboard (filter testid=$NAME)"
echo "Logs:   kubectl -n loadtest logs -f job/$NAME"
echo "Stop:   kubectl -n loadtest delete job $NAME"