#!/bin/bash

# Wave 78: gRPC Load Testing Script for Foxhunt HFT Services
# Tests Trading Service, Backtesting Service, ML Training Service, and API Gateway
# Uses ghz (https://ghz.sh) for gRPC load testing

set -e

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

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
FOXHUNT_ROOT="$(dirname "$SCRIPT_DIR")"
RESULTS_DIR="$FOXHUNT_ROOT/test_results/grpc_load_tests"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)

# Ensure PATH includes Go binaries
export PATH="$HOME/go/bin:$HOME/go-install/go/bin:$PATH"

# Service configurations
TRADING_SERVICE_HOST="localhost:50051"
BACKTESTING_SERVICE_HOST="localhost:50052"
ML_TRAINING_SERVICE_HOST="localhost:50053"
API_GATEWAY_HOST="localhost:50050"

# Proto file locations
TRADING_PROTO="$FOXHUNT_ROOT/services/trading_service/proto/trading.proto"
ML_TRAINING_PROTO="$FOXHUNT_ROOT/services/ml_training_service/proto/ml_training.proto"
HEALTH_PROTO="$FOXHUNT_ROOT/tli/proto/health.proto"

# Test parameters (configurable)
DURATION="${DURATION:-10s}"          # Test duration
RPS="${RPS:-100}"                    # Requests per second
CONNECTIONS="${CONNECTIONS:-10}"     # Number of connections
CONCURRENCY="${CONCURRENCY:-50}"     # Number of concurrent workers

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Foxhunt gRPC Load Testing - Wave 78${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

# Create results directory
mkdir -p "$RESULTS_DIR"

# Function to check if service is running
check_service() {
    local host=$1
    local service_name=$2

    echo -e "${YELLOW}Checking $service_name at $host...${NC}"

    if nc -z -w 5 $(echo $host | cut -d: -f1) $(echo $host | cut -d: -f2) 2>/dev/null; then
        echo -e "${GREEN}✓ $service_name is running${NC}"
        return 0
    else
        echo -e "${RED}✗ $service_name is not running${NC}"
        return 1
    fi
}

# Function to run ghz test
run_ghz_test() {
    local service_name=$1
    local host=$2
    local proto_file=$3
    local method=$4
    local data=$5
    local output_file="$RESULTS_DIR/${service_name}_${TIMESTAMP}.json"

    echo ""
    echo -e "${BLUE}Testing $service_name - $method${NC}"
    echo -e "${YELLOW}Target: $host${NC}"
    echo -e "${YELLOW}Duration: $DURATION | RPS: $RPS | Connections: $CONNECTIONS | Concurrency: $CONCURRENCY${NC}"

    # Run ghz with comprehensive metrics
    $HOME/go/bin/ghz \
        --insecure \
        --proto "$proto_file" \
        --call "$method" \
        -d "$data" \
        --duration "$DURATION" \
        --rps "$RPS" \
        --connections "$CONNECTIONS" \
        --concurrency "$CONCURRENCY" \
        --format json \
        --output "$output_file" \
        "$host" \
        2>&1 | tee >(cat >&2)

    # Also save human-readable output
    $HOME/go/bin/ghz \
        --insecure \
        --proto "$proto_file" \
        --call "$method" \
        -d "$data" \
        --duration "$DURATION" \
        --rps "$RPS" \
        --connections "$CONNECTIONS" \
        --concurrency "$CONCURRENCY" \
        --format pretty \
        "$host" \
        > "$RESULTS_DIR/${service_name}_${TIMESTAMP}.txt"

    echo -e "${GREEN}✓ Results saved to $output_file${NC}"

    # Display summary
    if [ -f "$output_file" ]; then
        echo ""
        echo -e "${BLUE}Summary:${NC}"
        cat "$output_file" | jq -r '
            "Total requests: " + (.count | tostring) +
            "\nAverage latency: " + (.average | tostring) + "ms" +
            "\nP50 latency: " + (.latencies.P50 | tostring) + "ms" +
            "\nP95 latency: " + (.latencies.P95 | tostring) + "ms" +
            "\nP99 latency: " + (.latencies.P99 | tostring) + "ms" +
            "\nMax latency: " + (.latencies.max | tostring) + "ms" +
            "\nRPS: " + (.rps | tostring) +
            "\nErrors: " + (if .errorDistribution then (.errorDistribution | length | tostring) else "0" end)
        ' 2>/dev/null || echo "Summary parsing failed (check JSON file directly)"
    fi
}

# Function to test Trading Service
test_trading_service() {
    echo ""
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}Testing Trading Service (Port 50051)${NC}"
    echo -e "${BLUE}========================================${NC}"

    if ! check_service "$TRADING_SERVICE_HOST" "Trading Service"; then
        echo -e "${RED}Skipping Trading Service tests${NC}"
        return 1
    fi

    # Test 1: GetOrderStatus (simple unary RPC)
    run_ghz_test \
        "trading_service_order_status" \
        "$TRADING_SERVICE_HOST" \
        "$TRADING_PROTO" \
        "trading.TradingService.GetOrderStatus" \
        '{"order_id": "test-order-123"}'

    # Test 2: GetPositions (unary RPC with optional fields)
    run_ghz_test \
        "trading_service_positions" \
        "$TRADING_SERVICE_HOST" \
        "$TRADING_PROTO" \
        "trading.TradingService.GetPositions" \
        '{"account_id": "test-account", "symbol": "AAPL"}'

    # Test 3: GetPortfolioSummary
    run_ghz_test \
        "trading_service_portfolio" \
        "$TRADING_SERVICE_HOST" \
        "$TRADING_PROTO" \
        "trading.TradingService.GetPortfolioSummary" \
        '{"account_id": "test-account"}'

    # Test 4: GetOrderBook
    run_ghz_test \
        "trading_service_orderbook" \
        "$TRADING_SERVICE_HOST" \
        "$TRADING_PROTO" \
        "trading.TradingService.GetOrderBook" \
        '{"symbol": "AAPL", "depth": 10}'
}

# Function to test ML Training Service
test_ml_training_service() {
    echo ""
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}Testing ML Training Service (Port 50053)${NC}"
    echo -e "${BLUE}========================================${NC}"

    if ! check_service "$ML_TRAINING_SERVICE_HOST" "ML Training Service"; then
        echo -e "${RED}Skipping ML Training Service tests${NC}"
        return 1
    fi

    # Test 1: HealthCheck
    run_ghz_test \
        "ml_training_service_health" \
        "$ML_TRAINING_SERVICE_HOST" \
        "$ML_TRAINING_PROTO" \
        "ml_training.MLTrainingService.HealthCheck" \
        '{}'

    # Test 2: ListAvailableModels
    run_ghz_test \
        "ml_training_service_models" \
        "$ML_TRAINING_SERVICE_HOST" \
        "$ML_TRAINING_PROTO" \
        "ml_training.MLTrainingService.ListAvailableModels" \
        '{}'

    # Test 3: ListTrainingJobs (with pagination)
    run_ghz_test \
        "ml_training_service_jobs" \
        "$ML_TRAINING_SERVICE_HOST" \
        "$ML_TRAINING_PROTO" \
        "ml_training.MLTrainingService.ListTrainingJobs" \
        '{"page": 1, "page_size": 10}'
}

# Function to test Health endpoints
test_health_endpoints() {
    echo ""
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}Testing Health Endpoints (All Services)${NC}"
    echo -e "${BLUE}========================================${NC}"

    # Trading Service Health
    if check_service "$TRADING_SERVICE_HOST" "Trading Service"; then
        run_ghz_test \
            "health_trading" \
            "$TRADING_SERVICE_HOST" \
            "$HEALTH_PROTO" \
            "grpc.health.v1.Health.Check" \
            '{"service": ""}'
    fi

    # ML Training Service Health (already has HealthCheck in its proto)
    # Covered above

    # API Gateway Health
    if check_service "$API_GATEWAY_HOST" "API Gateway"; then
        run_ghz_test \
            "health_gateway" \
            "$API_GATEWAY_HOST" \
            "$HEALTH_PROTO" \
            "grpc.health.v1.Health.Check" \
            '{"service": ""}'
    fi
}

# High-load stress test
stress_test_trading_service() {
    echo ""
    echo -e "${BLUE}========================================${NC}"
    echo -e "${BLUE}STRESS TEST: Trading Service${NC}"
    echo -e "${BLUE}========================================${NC}"

    if ! check_service "$TRADING_SERVICE_HOST" "Trading Service"; then
        echo -e "${RED}Skipping stress test${NC}"
        return 1
    fi

    echo -e "${YELLOW}Running high-load test: 1000 RPS, 100 connections, 30s duration${NC}"

    # Override parameters for stress test
    $HOME/go/bin/ghz \
        --insecure \
        --proto "$TRADING_PROTO" \
        --call "trading.TradingService.GetOrderStatus" \
        -d '{"order_id": "stress-test-order"}' \
        --duration "30s" \
        --rps 1000 \
        --connections 100 \
        --concurrency 200 \
        --format json \
        --output "$RESULTS_DIR/stress_test_trading_${TIMESTAMP}.json" \
        "$TRADING_SERVICE_HOST" \
        2>&1 | tee >(cat >&2)

    $HOME/go/bin/ghz \
        --insecure \
        --proto "$TRADING_PROTO" \
        --call "trading.TradingService.GetOrderStatus" \
        -d '{"order_id": "stress-test-order"}' \
        --duration "30s" \
        --rps 1000 \
        --connections 100 \
        --concurrency 200 \
        --format pretty \
        "$TRADING_SERVICE_HOST" \
        > "$RESULTS_DIR/stress_test_trading_${TIMESTAMP}.txt"

    echo -e "${GREEN}✓ Stress test results saved${NC}"
}

# Main execution
main() {
    local test_mode="${1:-all}"

    echo -e "${BLUE}Test mode: $test_mode${NC}"
    echo ""

    case "$test_mode" in
        trading)
            test_trading_service
            ;;
        ml)
            test_ml_training_service
            ;;
        health)
            test_health_endpoints
            ;;
        stress)
            stress_test_trading_service
            ;;
        all)
            test_trading_service
            test_ml_training_service
            test_health_endpoints
            ;;
        *)
            echo -e "${RED}Unknown test mode: $test_mode${NC}"
            echo "Usage: $0 [trading|ml|health|stress|all]"
            exit 1
            ;;
    esac

    echo ""
    echo -e "${GREEN}========================================${NC}"
    echo -e "${GREEN}All tests completed!${NC}"
    echo -e "${GREEN}Results saved to: $RESULTS_DIR${NC}"
    echo -e "${GREEN}========================================${NC}"
    echo ""
    echo -e "${YELLOW}To view detailed results:${NC}"
    echo "  - JSON: cat $RESULTS_DIR/*_${TIMESTAMP}.json | jq ."
    echo "  - Text: cat $RESULTS_DIR/*_${TIMESTAMP}.txt"
    echo ""
    echo -e "${YELLOW}To run stress test:${NC}"
    echo "  $0 stress"
}

# Run main function with arguments
main "$@"