#!/bin/bash

# Resource Monitoring Suite for C-Relay
# Monitors memory and CPU usage during testing

set -e

# Configuration
RELAY_HOST="127.0.0.1"
RELAY_PORT="8888"
MONITOR_DURATION=60  # seconds
SAMPLE_INTERVAL=2    # seconds

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

# Metrics storage
CPU_SAMPLES=()
MEM_SAMPLES=()
CONNECTION_SAMPLES=()
TIMESTAMP_SAMPLES=()

# Function to get relay process info
get_relay_info() {
    local pid
    pid=$(pgrep -f "c_relay" | head -1)

    if [[ -z "$pid" ]]; then
        echo "0:0:0:0"
        return
    fi

    # Get CPU, memory, and other stats
    local ps_output
    ps_output=$(ps -p "$pid" -o pcpu,pmem,vsz,rss --no-headers 2>/dev/null || echo "0.0 0.0 0 0")

    # Get connection count
    local connections
    connections=$(netstat -t 2>/dev/null | grep ":$RELAY_PORT" | wc -l 2>/dev/null || echo "0")

    echo "$ps_output $connections"
}

# Function to monitor resources
monitor_resources() {
    local duration="$1"
    local interval="$2"

    echo "=========================================="
    echo "Resource Monitoring Started"
    echo "=========================================="
    echo "Duration: ${duration}s, Interval: ${interval}s"
    echo ""

    # Clear arrays
    CPU_SAMPLES=()
    MEM_SAMPLES=()
    CONNECTION_SAMPLES=()
    TIMESTAMP_SAMPLES=()

    local start_time
    start_time=$(date +%s)
    local sample_count=0

    echo "Time | CPU% | Mem% | VSZ(KB) | RSS(KB) | Connections"
    echo "-----+------+------+---------+---------+------------"

    while [[ $(($(date +%s) - start_time)) -lt duration ]]; do
        local relay_info
        relay_info=$(get_relay_info)

        if [[ "$relay_info" != "0:0:0:0" ]]; then
            local cpu mem vsz rss connections
            IFS=' ' read -r cpu mem vsz rss connections <<< "$relay_info"

            # Store samples
            CPU_SAMPLES+=("$cpu")
            MEM_SAMPLES+=("$mem")
            CONNECTION_SAMPLES+=("$connections")
            TIMESTAMP_SAMPLES+=("$sample_count")

            # Display current stats
            local elapsed
            elapsed=$(($(date +%s) - start_time))
            printf "%4ds | %4.1f | %4.1f | %7s | %7s | %10s\n" \
                   "$elapsed" "$cpu" "$mem" "$vsz" "$rss" "$connections"
        else
            echo "  -- | Relay process not found --"
        fi

        ((sample_count++))
        sleep "$interval"
    done

    echo ""
}

# Function to calculate statistics
calculate_stats() {
    local array_name="$1"
    local -n array_ref="$array_name"

    if [[ ${#array_ref[@]} -eq 0 ]]; then
        echo "0:0:0:0:0"
        return
    fi

    local sum=0
    local min=${array_ref[0]}
    local max=${array_ref[0]}

    for value in "${array_ref[@]}"; do
        # Use awk for floating point arithmetic
        sum=$(awk "BEGIN {print $sum + $value}")
        min=$(awk "BEGIN {print ($value < $min) ? $value : $min}")
        max=$(awk "BEGIN {print ($value > $max) ? $value : $max}")
    done

    local avg
    avg=$(awk "BEGIN {print $sum / ${#array_ref[@]} }")

    echo "$avg:$min:$max:$sum:${#array_ref[@]}"
}

# Function to generate resource report
generate_resource_report() {
    echo "=========================================="
    echo "Resource Monitoring Report"
    echo "=========================================="

    if [[ ${#CPU_SAMPLES[@]} -eq 0 ]]; then
        echo "No resource samples collected. Is the relay running?"
        return
    fi

    # Calculate statistics
    local cpu_stats mem_stats conn_stats
    cpu_stats=$(calculate_stats CPU_SAMPLES)
    mem_stats=$(calculate_stats MEM_SAMPLES)
    conn_stats=$(calculate_stats CONNECTION_SAMPLES)

    # Parse statistics
    IFS=':' read -r cpu_avg cpu_min cpu_max cpu_sum cpu_count <<< "$cpu_stats"
    IFS=':' read -r mem_avg mem_min mem_max mem_sum mem_count <<< "$mem_stats"
    IFS=':' read -r conn_avg conn_min conn_max conn_sum conn_count <<< "$conn_stats"

    echo "CPU Usage Statistics:"
    printf "  Average: %.2f%%\n" "$cpu_avg"
    printf "  Minimum: %.2f%%\n" "$cpu_min"
    printf "  Maximum: %.2f%%\n" "$cpu_max"
    printf "  Samples: %d\n" "$cpu_count"
    echo ""

    echo "Memory Usage Statistics:"
    printf "  Average: %.2f%%\n" "$mem_avg"
    printf "  Minimum: %.2f%%\n" "$mem_min"
    printf "  Maximum: %.2f%%\n" "$mem_max"
    printf "  Samples: %d\n" "$mem_count"
    echo ""

    echo "Connection Statistics:"
    printf "  Average: %.1f connections\n" "$conn_avg"
    printf "  Minimum: %.1f connections\n" "$conn_min"
    printf "  Maximum: %.1f connections\n" "$conn_max"
    printf "  Samples: %d\n" "$conn_count"
    echo ""

    # Performance assessment
    echo "Performance Assessment:"
    if awk "BEGIN {exit !($cpu_avg < 50)}"; then
        echo -e "  ${GREEN}✓ CPU usage is acceptable${NC}"
    else
        echo -e "  ${RED}✗ CPU usage is high${NC}"
    fi

    if awk "BEGIN {exit !($mem_avg < 80)}"; then
        echo -e "  ${GREEN}✓ Memory usage is acceptable${NC}"
    else
        echo -e "  ${RED}✗ Memory usage is high${NC}"
    fi

    if [[ $(awk "BEGIN {print int($conn_max)}") -gt 0 ]]; then
        echo -e "  ${GREEN}✓ Relay is handling connections${NC}"
    else
        echo -e "  ${YELLOW}⚠ No active connections detected${NC}"
    fi
}

# Function to run load test with monitoring
run_monitored_load_test() {
    local test_name="$1"
    local description="$2"

    echo "=========================================="
    echo "Monitored Load Test: $test_name"
    echo "=========================================="
    echo "Description: $description"
    echo ""

    # Start monitoring in background
    monitor_resources 30 2 &
    local monitor_pid=$!

    # Wait a moment for monitoring to start
    sleep 2

    # Run a simple load test (create multiple subscriptions)
    echo "Running load test..."
    for i in {1..20}; do
        echo "[\"REQ\",\"monitor_test_${i}\",{}]" | timeout 3 websocat -B 1048576 ws://$RELAY_HOST:$RELAY_PORT >/dev/null 2>&1 &
    done

    # Let the load run for a bit
    sleep 10

    # Clean up subscriptions
    echo "Cleaning up test subscriptions..."
    for i in {1..20}; do
        echo "[\"CLOSE\",\"monitor_test_${i}\"]" | timeout 3 websocat -B 1048576 ws://$RELAY_HOST:$RELAY_PORT >/dev/null 2>&1 &
    done

    # Wait for monitoring to complete
    sleep 5
    kill "$monitor_pid" 2>/dev/null || true
    wait "$monitor_pid" 2>/dev/null || true

    echo ""
}

echo "=========================================="
echo "C-Relay Resource Monitoring Suite"
echo "=========================================="
echo "Monitoring relay at ws://$RELAY_HOST:$RELAY_PORT"
echo ""

# Check if relay is running
if ! pgrep -f "c_relay" >/dev/null 2>&1; then
    echo -e "${RED}Relay process not found. Please start the relay first.${NC}"
    echo "Use: ./make_and_restart_relay.sh"
    exit 1
fi

echo -e "${GREEN}✓ Relay process found${NC}"
echo ""

# Run baseline monitoring
echo "=== Baseline Resource Monitoring ==="
monitor_resources 15 2
generate_resource_report
echo ""

# Run monitored load test
run_monitored_load_test "Subscription Load Test" "Creating and closing multiple subscriptions while monitoring resources"
generate_resource_report
echo ""

echo "=========================================="
echo "Resource Monitoring Complete"
echo "=========================================="
echo "Resource monitoring completed. Review the statistics above."
echo "High CPU/memory usage may indicate performance issues."