#!/bin/bash

# Configuration
NUM_PROCESSES=$1
BATCH_SIZE=$2
WAIT_TIME_MS=$3
duration_string=$4

# Parse duration_string into parts, replacing empty fields with 0
read -ra duration_parts <<< "$(echo "$duration_string" | tr ',' ' ')"

# Set duration variables with defaults (0 if unset or empty)
DURATION_DAYS=${duration_parts[0]:-0}
DURATION_HOURS=${duration_parts[1]:-0}
DURATION_MINUTES=${duration_parts[2]:-0}

# Convert to integers to ensure numeric values
DURATION_DAYS=${DURATION_DAYS:-0}
DURATION_HOURS=${DURATION_HOURS:-0}
DURATION_MINUTES=${DURATION_MINUTES:-0}

WRAPPER_SCRIPT="/home/ssm-user/test_harness/src/newversion/run_k6_wrapper.sh"
AGGREGATED_CSV="/home/ssm-user/test_harness/src/newversion/metrics/aggregated_metrics.csv"
S3_BUCKET="k6testharnessrepo"
ERROR_LOG="/home/ssm-user/test_harness/src/newversion/log/errors_aggregator.log"

# Function to start a K6 process using the wrapper
start_k6_process() {
    local process_id=$1
    echo "Starting K6 process $process_id at $(date)"
    nohup "$WRAPPER_SCRIPT" $BATCH_SIZE $WAIT_TIME_MS $DURATION_DAYS $DURATION_HOURS $DURATION_MINUTES $process_id &
    echo $! > "/home/ssm-user/test_harness/src/newversion/pid/k6_process_$process_id.pid"
}

# Function to check if a process is running
is_process_running() {
    local pid_file="$1"
    if [ -f "$pid_file" ]; then
        local pid=$(cat "$pid_file")
        if ps -p $pid > /dev/null 2>&1; then
            return 0 # Process is running
        fi
    fi
    return 1 # Process is not running
}

# Function to combine metrics
combine_metrics() {
    echo "Combining metrics from all processes at $(date)" >> "$ERROR_LOG"
    
    # Initialize aggregated CSV
    echo "Timestamp,ResponseTime(ms),ErrorRate" > "$AGGREGATED_CSV"
    
    # Find all process-specific CSV files in S3 and locally
    local_files="/home/ssm-user/test_harness/src/newversion/metrics/metrics_*.csv"
    aws s3 ls "s3://$S3_BUCKET/" | grep "metrics_.*\.csv" | awk '{print $4}' > /tmp/s3_files.txt

    # Download S3 files temporarily
    while IFS= read -r s3_file; do
        local_temp_file="/tmp/$s3_file"
        aws s3 cp "s3://$S3_BUCKET/$s3_file" "$local_temp_file" 2>>"$ERROR_LOG"
        if [ -f "$local_temp_file" ]; then
            tail -n +2 "$local_temp_file" >> "$AGGREGATED_CSV" # Skip header
            rm "$local_temp_file"
        fi
    done < /tmp/s3_files.txt
    rm /tmp/s3_files.txt

    # Append local files (skip headers)
    for local_file in $local_files; do
        if [ -f "$local_file" ]; then
            tail -n +2 "$local_file" >> "$AGGREGATED_CSV"
        fi
    done

    # Sort by timestamp and remove duplicates (if any)
    temp_file="/tmp/aggregated_metrics_temp.csv"
    sort -t',' -k1 "$AGGREGATED_CSV" | uniq > "$temp_file"
    echo "Timestamp,ResponseTime(ms),ErrorRate" > "$AGGREGATED_CSV"
    cat "$temp_file" >> "$AGGREGATED_CSV"
    rm "$temp_file"

    # Archive aggregated CSV to S3
    if [ -s "$AGGREGATED_CSV" ]; then
        aws s3 cp "$AGGREGATED_CSV" "s3://$S3_BUCKET/aggregated_metrics.csv" >> "$ERROR_LOG" 2>&1 || echo "Failed to upload aggregated CSV" >> "$ERROR_LOG"
    fi
}

# Start all processes with a 1-second stagger
for ((i=1; i<=NUM_PROCESSES; i++)); do
    start_k6_process $i
    if [ $i -lt $NUM_PROCESSES ]; then
        sleep 1
    fi
done

echo "All $NUM_PROCESSES K6 processes started."

# Wait for all processes to finish
all_finished=false
while [ "$all_finished" != "true" ]; do
    all_finished=true
    for ((i=1; i<=NUM_PROCESSES; i++)); do
        PID_FILE="/home/ssm-user/test_harness/src/newversion/pid/k6_process_$i.pid"
        if is_process_running "$PID_FILE"; then
            all_finished=false
            break
        fi
    done
    sleep 10 # Check every 10 seconds
done

# Run aggregation after all processes finish
combine_metrics

# Clean up PID files
for ((i=1; i<=NUM_PROCESSES; i++)); do
    PID_FILE="/home/ssm-user/test_harness/src/newversion/pid/k6_process_$i.pid"
    rm -f "$PID_FILE"
done

echo "All processes completed and metrics aggregated at $(date)"