#!/bin/bash

# Benchmark runner script for Gotify WebSocket performance testing

set -e

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

# Configuration
COMPOSE_FILE="docker-compose.benchmark.yml"
INSTANCE_PORTS=(8080 8081 8082 8083 8084)
INSTANCE_NAMES=("64" "128" "256" "512" "1024")

# Function to print colored output
print_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Function to check if services are healthy
wait_for_services() {
    print_info "Waiting for Gotify instances to be healthy..."
    
    for port in "${INSTANCE_PORTS[@]}"; do
        local max_attempts=30
        local attempt=0
        
        while [ $attempt -lt $max_attempts ]; do
            if curl -sf "http://localhost:${port}/health" > /dev/null 2>&1; then
                print_info "Instance on port ${port} is healthy"
                break
            fi
            
            attempt=$((attempt + 1))
            sleep 2
        done
        
        if [ $attempt -eq $max_attempts ]; then
            print_error "Instance on port ${port} failed to become healthy"
            return 1
        fi
    done
    
    print_info "All instances are healthy"
}

# Function to run benchmark against a specific instance
run_benchmark() {
    local instance_name=$1
    local port=$2
    local test_script=${3:-"websocket-simple.js"}
    
    print_info "Running benchmark against instance with ${instance_name} shards (port ${port})..."
    
    docker run --rm \
        --network gotify_benchmark-net \
        -v "$(pwd)/benchmark/k6:/scripts" \
        -e BASE_URL="http://gotify-${instance_name}:80" \
        grafana/k6:latest \
        run /scripts/${test_script}
}

# Function to compare all instances
compare_all_instances() {
    print_info "Running comparison benchmark across all instances..."
    
    for i in "${!INSTANCE_NAMES[@]}"; do
        local name="${INSTANCE_NAMES[$i]}"
        local port="${INSTANCE_PORTS[$i]}"
        
        print_info "Testing instance with ${name} shards..."
        run_benchmark "${name}" "${port}" "websocket-simple.js" > "benchmark/results/${name}-shards.log" 2>&1 || true
        sleep 5
    done
}

# Main execution
main() {
    print_info "Starting Gotify WebSocket Benchmark Suite"
    
    # Create results directory
    mkdir -p benchmark/results
    
    # Start services
    print_info "Starting Docker Compose services..."
    docker-compose -f ${COMPOSE_FILE} up -d --build
    
    # Wait for services to be ready
    wait_for_services
    
    # Give services a moment to fully initialize
    sleep 5
    
    # Parse command line arguments
    case "${1:-all}" in
        "all")
            compare_all_instances
            ;;
        "64"|"128"|"256"|"512"|"1024")
            local port_index=0
            for i in "${!INSTANCE_NAMES[@]}"; do
                if [ "${INSTANCE_NAMES[$i]}" = "$1" ]; then
                    port_index=$i
                    break
                fi
            done
            run_benchmark "$1" "${INSTANCE_PORTS[$port_index]}" "${2:-websocket-simple.js}"
            ;;
        "scale")
            local scale=${2:-"1k"}
            print_info "Running connection scaling test with ${scale} connections..."
            docker run --rm \
                --network gotify_benchmark-net \
                -v "$(pwd)/benchmark/k6:/scripts" \
                -e BASE_URL="http://gotify-256:80" \
                -e SCALE="${scale}" \
                grafana/k6:latest \
                run /scripts/connection-scaling.js
            ;;
        "stop")
            print_info "Stopping Docker Compose services..."
            docker-compose -f ${COMPOSE_FILE} down -v
            ;;
        *)
            echo "Usage: $0 [all|64|128|256|512|1024|scale|stop] [test-script|scale-size]"
            echo ""
            echo "Commands:"
            echo "  all              - Run benchmarks against all instances"
            echo "  64|128|256|512|1024 - Run benchmark against specific shard count"
            echo "  scale [1k|10k|100k] - Run connection scaling test"
            echo "  stop              - Stop all services"
            exit 1
            ;;
    esac
    
    print_info "Benchmark completed. Check benchmark/results/ for detailed logs."
}

main "$@"