#!/bin/bash
# gRPC Load Testing Script
# Requires: ghz (go install github.com/bojand/ghz/cmd/ghz@latest)
#
# This script executes comprehensive gRPC load tests against Foxhunt services

set -euo pipefail

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

# Configuration
TRADING_ENDPOINT="${GRPC_ENDPOINT:-localhost:50051}"
PROTO_PATH="${PROTO_PATH:-tli/proto/trading.proto}"
JWT_SECRET="${JWT_SECRET:-test-secret-key-for-load-testing}"
OUTPUT_DIR="${OUTPUT_DIR:-./load_test_results}"

# Performance targets (from Wave 76 validation)
TARGET_P99_LATENCY_US=10000  # 10μs = 10,000ns
TARGET_THROUGHPUT_RPS=100000  # 100K req/s
TARGET_ERROR_RATE_PCT=0.1

# Create output directory
mkdir -p "$OUTPUT_DIR"

# Function to check if ghz is installed
check_ghz() {
    if ! command -v ghz &> /dev/null; then
        echo -e "${RED}ERROR: ghz is not installed${NC}"
        echo -e "${YELLOW}Install with: go install github.com/bojand/ghz/cmd/ghz@latest${NC}"
        echo -e "${YELLOW}Or use Docker: docker pull ghcr.io/bojand/ghz:latest${NC}"
        exit 1
    fi
}

# Function to check if service is running
check_service() {
    echo -e "${BLUE}Checking if service is running on ${TRADING_ENDPOINT}...${NC}"

    if ! nc -z localhost 50051 2>/dev/null; then
        echo -e "${RED}ERROR: Trading service not running on ${TRADING_ENDPOINT}${NC}"
        echo -e "${YELLOW}Start with: ./target/release/trading_service${NC}"
        exit 1
    fi

    echo -e "${GREEN}✓ Service is running${NC}"
}

# Function to generate JWT token
generate_jwt() {
    # This is a simplified example - production would use proper JWT library
    echo "test-jwt-token-placeholder"
}

# Function to run normal load test
run_normal_load_test() {
    echo -e "\n${BLUE}=== NORMAL LOAD TEST ===${NC}"
    echo -e "${BLUE}Clients: 1,000 | Duration: 60s${NC}\n"

    local output_file="$OUTPUT_DIR/normal_load_$(date +%Y%m%d_%H%M%S).json"
    local jwt_token=$(generate_jwt)

    ghz --insecure \
        --proto="$PROTO_PATH" \
        --call="trading.TradingService/GetPositions" \
        --connections=1000 \
        --concurrency=1000 \
        --duration=60s \
        --rps=0 \
        --data='{"account_id":"test-account"}' \
        --metadata="{\"authorization\":\"Bearer $jwt_token\"}" \
        --format=json \
        --output="$output_file" \
        "$TRADING_ENDPOINT"

    echo -e "${GREEN}✓ Normal load test complete${NC}"
    echo -e "Results saved to: $output_file"

    # Parse and validate results
    validate_results "$output_file" "Normal Load"
}

# Function to run spike load test
run_spike_load_test() {
    echo -e "\n${BLUE}=== SPIKE LOAD TEST ===${NC}"
    echo -e "${BLUE}Ramp: 0→10,000 clients | Duration: 30s${NC}\n"

    local output_file="$OUTPUT_DIR/spike_load_$(date +%Y%m%d_%H%M%S).json"
    local jwt_token=$(generate_jwt)

    ghz --insecure \
        --proto="$PROTO_PATH" \
        --call="trading.TradingService/GetPositions" \
        --connections=10000 \
        --concurrency=10000 \
        --duration=30s \
        --rps=0 \
        --data='{"account_id":"test-account"}' \
        --metadata="{\"authorization\":\"Bearer $jwt_token\"}" \
        --format=json \
        --output="$output_file" \
        "$TRADING_ENDPOINT"

    echo -e "${GREEN}✓ Spike load test complete${NC}"
    echo -e "Results saved to: $output_file"

    # Parse and validate results
    validate_results "$output_file" "Spike Load"
}

# Function to run sustained load test
run_sustained_load_test() {
    echo -e "\n${BLUE}=== SUSTAINED LOAD TEST ===${NC}"
    echo -e "${BLUE}Clients: 100 | Duration: 300s (5 minutes)${NC}\n"

    local output_file="$OUTPUT_DIR/sustained_load_$(date +%Y%m%d_%H%M%S).json"
    local jwt_token=$(generate_jwt)

    ghz --insecure \
        --proto="$PROTO_PATH" \
        --call="trading.TradingService/GetPositions" \
        --connections=100 \
        --concurrency=100 \
        --duration=300s \
        --rps=0 \
        --data='{"account_id":"test-account"}' \
        --metadata="{\"authorization\":\"Bearer $jwt_token\"}" \
        --format=json \
        --output="$output_file" \
        "$TRADING_ENDPOINT"

    echo -e "${GREEN}✓ Sustained load test complete${NC}"
    echo -e "Results saved to: $output_file"

    # Parse and validate results
    validate_results "$output_file" "Sustained Load"
}

# Function to validate results against targets
validate_results() {
    local results_file="$1"
    local test_name="$2"

    echo -e "\n${BLUE}Validating $test_name results...${NC}"

    if [ ! -f "$results_file" ]; then
        echo -e "${RED}✗ Results file not found${NC}"
        return 1
    fi

    # Parse JSON results using jq
    if command -v jq &> /dev/null; then
        local p50=$(jq -r '.latencies.p50' "$results_file" 2>/dev/null || echo "0")
        local p95=$(jq -r '.latencies.p95' "$results_file" 2>/dev/null || echo "0")
        local p99=$(jq -r '.latencies.p99' "$results_file" 2>/dev/null || echo "0")
        local rps=$(jq -r '.rps' "$results_file" 2>/dev/null || echo "0")
        local error_pct=$(jq -r '.errorRate' "$results_file" 2>/dev/null || echo "0")

        echo -e "\n${BLUE}Performance Metrics:${NC}"
        echo -e "  P50 Latency: ${p50}ns"
        echo -e "  P95 Latency: ${p95}ns"
        echo -e "  P99 Latency: ${p99}ns"
        echo -e "  Throughput:  ${rps} req/s"
        echo -e "  Error Rate:  ${error_pct}%"

        # Validate against targets
        echo -e "\n${BLUE}Target Validation:${NC}"

        # P99 Latency check
        if (( $(echo "$p99 < $TARGET_P99_LATENCY_US" | bc -l) )); then
            echo -e "  ${GREEN}✓ P99 Latency: ${p99}ns < ${TARGET_P99_LATENCY_US}ns target${NC}"
        else
            echo -e "  ${RED}✗ P99 Latency: ${p99}ns >= ${TARGET_P99_LATENCY_US}ns target${NC}"
        fi

        # Throughput check
        if (( $(echo "$rps > $TARGET_THROUGHPUT_RPS" | bc -l) )); then
            echo -e "  ${GREEN}✓ Throughput: ${rps} req/s > ${TARGET_THROUGHPUT_RPS} req/s target${NC}"
        else
            echo -e "  ${YELLOW}⚠ Throughput: ${rps} req/s <= ${TARGET_THROUGHPUT_RPS} req/s target${NC}"
        fi

        # Error rate check
        if (( $(echo "$error_pct < $TARGET_ERROR_RATE_PCT" | bc -l) )); then
            echo -e "  ${GREEN}✓ Error Rate: ${error_pct}% < ${TARGET_ERROR_RATE_PCT}% target${NC}"
        else
            echo -e "  ${RED}✗ Error Rate: ${error_pct}% >= ${TARGET_ERROR_RATE_PCT}% target${NC}"
        fi
    else
        echo -e "${YELLOW}⚠ jq not installed, skipping detailed validation${NC}"
        echo -e "  Install with: sudo apt-get install jq"
    fi
}

# Function to generate summary report
generate_summary_report() {
    echo -e "\n${BLUE}=== LOAD TEST SUMMARY ===${NC}\n"

    local report_file="$OUTPUT_DIR/summary_report_$(date +%Y%m%d_%H%M%S).md"

    cat > "$report_file" << EOF
# gRPC Load Test Summary Report

**Date**: $(date '+%Y-%m-%d %H:%M:%S')
**Endpoint**: $TRADING_ENDPOINT
**Proto**: $PROTO_PATH

## Test Results

### Normal Load Test (1K clients, 60s)
- Status: COMPLETED
- Results: $OUTPUT_DIR/normal_load_*.json

### Spike Load Test (10K clients, 30s)
- Status: COMPLETED
- Results: $OUTPUT_DIR/spike_load_*.json

### Sustained Load Test (100 clients, 5m)
- Status: COMPLETED
- Results: $OUTPUT_DIR/sustained_load_*.json

## Performance Targets

| Metric | Target | Status |
|--------|--------|--------|
| P99 Latency | <10μs | See JSON results |
| Throughput | >100K req/s | See JSON results |
| Error Rate | <0.1% | See JSON results |

## Files Generated

\`\`\`bash
ls -lh $OUTPUT_DIR/
\`\`\`

## Next Steps

1. Review detailed JSON results
2. Analyze latency histograms
3. Compare against baseline
4. Document any anomalies
5. Update capacity planning

---

*Generated by: scripts/grpc_load_test.sh*
EOF

    echo -e "${GREEN}✓ Summary report generated: $report_file${NC}"
    cat "$report_file"
}

# Main execution
main() {
    echo -e "${BLUE}╔══════════════════════════════════════════╗${NC}"
    echo -e "${BLUE}║  Foxhunt gRPC Load Testing Suite        ║${NC}"
    echo -e "${BLUE}╚══════════════════════════════════════════╝${NC}"

    # Check prerequisites
    check_ghz
    check_service

    # Run test scenarios
    case "${1:-all}" in
        normal)
            run_normal_load_test
            ;;
        spike)
            run_spike_load_test
            ;;
        sustained)
            run_sustained_load_test
            ;;
        all)
            run_normal_load_test
            echo -e "\n${BLUE}Waiting 30s before next test...${NC}"
            sleep 30

            run_spike_load_test
            echo -e "\n${BLUE}Waiting 30s before next test...${NC}"
            sleep 30

            run_sustained_load_test

            # Generate summary
            generate_summary_report
            ;;
        *)
            echo "Usage: $0 {normal|spike|sustained|all}"
            echo ""
            echo "Tests:"
            echo "  normal    - 1K clients for 60s"
            echo "  spike     - 0→10K ramp for 30s"
            echo "  sustained - 100 clients for 5m"
            echo "  all       - Run all tests sequentially"
            exit 1
            ;;
    esac

    echo -e "\n${GREEN}╔══════════════════════════════════════════╗${NC}"
    echo -e "${GREEN}║  Load Testing Complete                   ║${NC}"
    echo -e "${GREEN}╚══════════════════════════════════════════╝${NC}"
}

# Run main with all arguments
main "$@"