#!/bin/bash
# Wave 79 Load Testing Script
# After GetOrderStatus issue is fixed and TLS is configured

set -euo pipefail

# Configuration
CERT_DIR="${CERT_DIR:-certs/production}"
RESULTS_DIR="${RESULTS_DIR:-load_test_results/wave79}"
PROTO_DIR="services/trading_service/proto"
HOST="${HOST:-localhost}"
PORT="${PORT:-50050}"

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Create results directory
mkdir -p "$RESULTS_DIR"

echo "================================================"
echo "Wave 79 Load Testing Suite"
echo "================================================"
echo "Host: $HOST:$PORT"
echo "Certificates: $CERT_DIR"
echo "Results: $RESULTS_DIR"
echo "================================================"
echo

# Check prerequisites
echo "Checking prerequisites..."

# Check if ghz is installed
if ! command -v ghz &> /dev/null; then
    echo -e "${RED}ERROR: ghz not found. Install with: go install github.com/bojand/ghz/cmd/ghz@latest${NC}"
    exit 1
fi

# Check if certificates exist
if [ ! -f "$CERT_DIR/ca/ca-cert.pem" ]; then
    echo -e "${RED}ERROR: CA certificate not found at $CERT_DIR/ca/ca-cert.pem${NC}"
    exit 1
fi

if [ ! -f "$CERT_DIR/foxhunt-cert.pem" ]; then
    echo -e "${RED}ERROR: Client certificate not found at $CERT_DIR/foxhunt-cert.pem${NC}"
    exit 1
fi

if [ ! -f "$CERT_DIR/foxhunt-key.pem" ]; then
    echo -e "${RED}ERROR: Client key not found at $CERT_DIR/foxhunt-key.pem${NC}"
    exit 1
fi

# Check if proto file exists
if [ ! -f "$PROTO_DIR/trading.proto" ]; then
    echo -e "${RED}ERROR: Proto file not found at $PROTO_DIR/trading.proto${NC}"
    exit 1
fi

echo -e "${GREEN}✓ All prerequisites met${NC}"
echo

# Common ghz flags
GHZ_FLAGS=(
    --cacert "$CERT_DIR/ca/ca-cert.pem"
    --cert "$CERT_DIR/foxhunt-cert.pem"
    --key "$CERT_DIR/foxhunt-key.pem"
    --proto "$PROTO_DIR/trading.proto"
    --call "trading.TradingService.GetOrderStatus"
    --data '{"order_id": "load_test_order_001"}'
    "$HOST:$PORT"
)

# Test 1: Normal Load (1K concurrent, 60s, 100K RPS target)
echo "================================================"
echo "Test 1: Normal Load"
echo "  Connections: 1,000"
echo "  Concurrency: 1,000"
echo "  Duration: 60s"
echo "  Target RPS: 100,000"
echo "================================================"
ghz "${GHZ_FLAGS[@]}" \
    --duration 60s \
    --rps 100000 \
    --connections 1000 \
    --concurrency 1000 \
    > "$RESULTS_DIR/test1_normal_load.json"

# Extract key metrics
echo -e "${GREEN}✓ Test 1 complete${NC}"
grep "Requests/sec:" "$RESULTS_DIR/test1_normal_load.json" || true
grep "Error distribution:" -A 10 "$RESULTS_DIR/test1_normal_load.json" || true
echo

# Test 2: Spike Load (0→10K ramp over 30s)
echo "================================================"
echo "Test 2: Spike Load (Ramp 0→10K)"
echo "  Initial connections: 100"
echo "  Max connections: 10,000"
echo "  Ramp duration: 30s"
echo "  Target RPS: 200,000"
echo "================================================"
# Note: ghz doesn't support native ramp, so we'll approximate with burst
ghz "${GHZ_FLAGS[@]}" \
    --duration 30s \
    --rps 200000 \
    --connections 10000 \
    --concurrency 10000 \
    > "$RESULTS_DIR/test2_spike_load.json"

echo -e "${GREEN}✓ Test 2 complete${NC}"
grep "Requests/sec:" "$RESULTS_DIR/test2_spike_load.json" || true
grep "Error distribution:" -A 10 "$RESULTS_DIR/test2_spike_load.json" || true
echo

# Test 3: Stress Test (10K concurrent, 60s)
echo "================================================"
echo "Test 3: Stress Test"
echo "  Connections: 10,000"
echo "  Concurrency: 10,000"
echo "  Duration: 60s"
echo "  Target RPS: 200,000"
echo "================================================"
ghz "${GHZ_FLAGS[@]}" \
    --duration 60s \
    --rps 200000 \
    --connections 10000 \
    --concurrency 10000 \
    > "$RESULTS_DIR/test3_stress_test.json"

echo -e "${GREEN}✓ Test 3 complete${NC}"
grep "Requests/sec:" "$RESULTS_DIR/test3_stress_test.json" || true
grep "Error distribution:" -A 10 "$RESULTS_DIR/test3_stress_test.json" || true
echo

# Test 4: Sustained Load (5K concurrent, 5 minutes)
echo "================================================"
echo "Test 4: Sustained Load"
echo "  Connections: 5,000"
echo "  Concurrency: 5,000"
echo "  Duration: 300s (5 minutes)"
echo "  Target RPS: 150,000"
echo "================================================"
ghz "${GHZ_FLAGS[@]}" \
    --duration 300s \
    --rps 150000 \
    --connections 5000 \
    --concurrency 5000 \
    > "$RESULTS_DIR/test4_sustained_load.json"

echo -e "${GREEN}✓ Test 4 complete${NC}"
grep "Requests/sec:" "$RESULTS_DIR/test4_sustained_load.json" || true
grep "Error distribution:" -A 10 "$RESULTS_DIR/test4_sustained_load.json" || true
echo

# Summary
echo "================================================"
echo "Test Suite Complete"
echo "================================================"
echo "Results saved to: $RESULTS_DIR"
echo
echo "Summary:"
echo "--------"
for i in 1 2 3 4; do
    file="$RESULTS_DIR/test${i}_*.json"
    if [ -f $file ]; then
        rps=$(grep "Requests/sec:" $file | awk '{print $2}')
        errors=$(grep -A 1 "Status code distribution:" $file | grep -v "Status code" | wc -l)
        echo "Test $i: $rps req/s (errors: $errors)"
    fi
done
echo
echo -e "${GREEN}All tests completed successfully!${NC}"
echo
echo "Compare to Wave 78 baseline: 211,000 req/s"
echo "Target: >200,000 req/s with <0.1% error rate"