#!/bin/bash

# Configuration
K6_SCRIPT="/home/ssm-user/test_harness/src/K6-Updated-MemEfficientS3Archive.js"
BATCH_SIZE=$1
WAIT_TIME_MS=$2
DURATION_DAYS=$3
DURATION_HOURS=$4
DURATION_MINUTES=$5
PROCESS_ID=$6
LOG_FILE="/home/ssm-user/test_harness/src/newversion/log/k6_process_${PROCESS_ID}.log"
JSON_OUTPUT="/home/ssm-user/test_harness/src/newversion/metrics/k6_metrics_${PROCESS_ID}.json"
CSV_FILE="/home/ssm-user/test_harness/src/newversion/metrics/metrics_${PROCESS_ID}.csv"
S3_BUCKET="k6testharnessrepo"
MAX_FILE_SIZE_BYTES=$((200 * 1024 * 1024)) # 200 MB
ERROR_LOG="/home/ssm-user/test_harness/src/newversion/log/errors_${PROCESS_ID}.log"

# Ensure directories exist
mkdir -p "$(dirname "$LOG_FILE")" "$(dirname "$CSV_FILE")" "$(dirname "$JSON_OUTPUT")"

# Function to get file size in bytes
get_file_size() {
    local file="$1"
    if [ -f "$file" ]; then
        stat -c %s "$file" 2>>"$ERROR_LOG" || echo 0
    else
        echo 0
    fi
}

# Function to archive CSV to S3
archive_to_s3() {
    local file="$1"
    local timestamp=$(date +%s)
    local archive_file="/home/ssm-user/test_harness/src/newversion/metrics/metrics_${PROCESS_ID}_${timestamp}.csv"
    
    if [ -f "$file" ]; then
        echo "Archiving $file (size: $(du -h "$file" | cut -f1)) to S3 at $(date)" >> "$LOG_FILE"
        if mv "$file" "$archive_file"; then
            echo "Renamed $file to $archive_file" >> "$LOG_FILE"
            if aws s3 cp "$archive_file" "s3://$S3_BUCKET/$(basename "$archive_file")" >> "$LOG_FILE" 2>>"$ERROR_LOG"; then
                echo "Successfully uploaded $archive_file to S3" >> "$LOG_FILE"
                rm "$archive_file" || echo "Failed to remove $archive_file" >> "$ERROR_LOG"
            else
                echo "S3 upload failed for $archive_file" >> "$ERROR_LOG"
                mv "$archive_file" "$file" # Restore on failure
            fi
        else
            echo "Failed to rename $file to $archive_file" >> "$ERROR_LOG"
        fi
    else
        echo "File $file does not exist, skipping archive" >> "$LOG_FILE"
    fi
}

# Function to process JSON to CSV
process_json_to_csv() {
    if [ -f "$JSON_OUTPUT" ] && [ -s "$JSON_OUTPUT" ]; then
        if command -v jq > /dev/null 2>&1; then
            # Initialize CSV file with header if needed
            if [ ! -f "$CSV_FILE" ]; then
                echo "Timestamp,ResponseTime(ms),ErrorRate" > "$CSV_FILE"
            fi

            # Extract Trend metrics (response_time)
            jq -r '.metrics | select(.type == "Trend" and .name == "response_time") | [.sample.timestamp, .values.avg, null] | join(",")' "$JSON_OUTPUT" >> "$CSV_FILE" 2>>"$ERROR_LOG"

            # Extract Rate metrics (error_rate)
            jq -r '.metrics | select(.type == "Rate" and .name == "error_rate") | [.sample.timestamp, null, .values.rate] | join(",")' "$JSON_OUTPUT" >> "$CSV_FILE" 2>>"$ERROR_LOG"

            # Log success
            echo "Processed metrics from $JSON_OUTPUT into $CSV_FILE" >> "$LOG_FILE"

            # Check CSV file size and archive
            file_size=$(get_file_size "$CSV_FILE")
            if [ "$file_size" -ge "$MAX_FILE_SIZE_BYTES" ]; then
                archive_to_s3 "$CSV_FILE"
                echo "Timestamp,ResponseTime(ms),ErrorRate" > "$CSV_FILE"
            fi
        else
            echo "jq not installed, skipping JSON processing" >> "$ERROR_LOG"
        fi
    fi
    # Clear JSON file after processing
    if [ -f "$JSON_OUTPUT" ]; then
        > "$JSON_OUTPUT"
    fi
}

# Calculate total duration in seconds
totalSeconds=$(( (DURATION_DAYS * 24 * 60 * 60) + (DURATION_HOURS * 60 * 60) + (DURATION_MINUTES * 60) ))
interval_seconds=60  # Run in 1-minute intervals
remaining_seconds=$total_seconds

# Run K6 in intervals
while [ $remaining_seconds -gt 0 ]; do
    if [ $remaining_seconds -ge $interval_seconds ]; then
        run_duration=$interval_seconds
    else
        run_duration=$remaining_seconds
    fi

    echo "Running k6 for $run_duration seconds (remaining: $remaining_seconds seconds)" >> "$LOG_FILE"

    # Run K6 for the current interval
    k6 run "$K6_SCRIPT" \
        --env BATCH_SIZE="$BATCH_SIZE" \
        --env WAIT_TIME_MS="$WAIT_TIME_MS" \
        --env DURATION_DAYS="$DURATION_DAYS" \
        --env DURATION_HOURS="$DURATION_HOURS" \
        --env DURATION_MINUTES="$((run_duration / 60))" \
        --env PROCESS_ID="$PROCESS_ID" \
        --out json="$JSON_OUTPUT" \
        --quiet \
        | tee -a "$LOG_FILE"

    # Process the JSON file after each interval
    echo "Processing JSON after interval at $(date)" >> "$LOG_FILE"
    process_json_to_csv

    remaining_seconds=$((remaining_seconds - run_duration))
done

# Archive any remaining CSV data
file_size=$(get_file_size "$CSV_FILE")
if [ "$file_size" -gt 0 ]; then
    archive_to_s3 "$CSV_FILE"
fi

# Final cleanup of JSON file
if [ -f "$JSON_OUTPUT" ]; then
    rm "$JSON_OUTPUT" || echo "Failed to remove $JSON_OUTPUT" >> "$ERROR_LOG"
fi

echo "K6 process $PROCESS_ID completed at $(date)" >> "$LOG_FILE"