#!/bin/bash

# ═══════════════════════════════════════════════════════════════════════════
# Server Capacity Metrics Capture
# Run this WHILE your load test is running
# ═══════════════════════════════════════════════════════════════════════════

PROCESS_NAME="${1:-broadcast_server}"
DURATION="${2:-60}"
INTERVAL="${3:-1}"
LOGFILE="server_metrics_$(date +%Y%m%d_%H%M%S).txt"
CSVFILE="server_metrics_$(date +%Y%m%d_%H%M%S).csv"

# Find the process
SERVER_PID=$(pgrep -f "$PROCESS_NAME" | head -1)

if [ -z "$SERVER_PID" ]; then
    echo "❌ Process '$PROCESS_NAME' not found"
    echo "Usage: $0 [process_name] [duration_seconds] [interval_seconds]"
    exit 1
fi

echo "═══════════════════════════════════════════════════════════════════"
echo "📊 Server Capacity Metrics Capture"
echo "═══════════════════════════════════════════════════════════════════"
echo "Process:    $PROCESS_NAME (PID: $SERVER_PID)"
echo "Duration:   ${DURATION}s"
echo "Interval:   ${INTERVAL}s"
echo "Log file:   $LOGFILE"
echo "CSV file:   $CSVFILE"
echo "═══════════════════════════════════════════════════════════════════"
echo ""

# Get baseline memory (before load)
BASELINE_RSS=$(ps -o rss= -p $SERVER_PID 2>/dev/null | tr -d ' ')
BASELINE_RSS_MB=$(echo "scale=2; ${BASELINE_RSS:-0} / 1024" | bc)

echo "Baseline RSS: ${BASELINE_RSS_MB} MB"
echo ""

# Write headers
{
    echo "═══════════════════════════════════════════════════════════════════"
    echo "SERVER CAPACITY METRICS - $(date)"
    echo "Process: $PROCESS_NAME (PID: $SERVER_PID)"
    echo "Baseline RSS: ${BASELINE_RSS_MB} MB"
    echo "═══════════════════════════════════════════════════════════════════"
    echo ""
} > "$LOGFILE"

echo "timestamp,elapsed_sec,cpu_pct,rss_mb,vsz_mb,threads,established,time_wait,fds,total_tcp" > "$CSVFILE"

# Tracking variables for peak/avg calculation
TOTAL_CPU=0
TOTAL_RSS=0
TOTAL_CONNS=0
PEAK_CPU=0
PEAK_RSS=0
PEAK_CONNS=0
SAMPLE_COUNT=0
START_TIME=$(date +%s)

echo "Capturing metrics for ${DURATION}s... (Ctrl+C to stop early)"
echo ""

capture_sample() {
    local ELAPSED=$1
    
    # Get CPU and memory from ps (more reliable than top for scripting)
    local PS_OUTPUT=$(ps -o %cpu=,rss=,vsz= -p $SERVER_PID 2>/dev/null)
    local CPU=$(echo "$PS_OUTPUT" | awk '{print $1}' | tr -d ' ')
    local RSS=$(echo "$PS_OUTPUT" | awk '{print $2}' | tr -d ' ')
    local VSZ=$(echo "$PS_OUTPUT" | awk '{print $3}' | tr -d ' ')
    
    # Convert to MB
    local RSS_MB=$(echo "scale=2; ${RSS:-0} / 1024" | bc)
    local VSZ_MB=$(echo "scale=2; ${VSZ:-0} / 1024" | bc)
    
    # Thread count
    local THREADS=$(ps -M $SERVER_PID 2>/dev/null | tail -n +2 | wc -l | tr -d ' ')
    
    # Network stats
    local ESTABLISHED=$(netstat -an 2>/dev/null | grep -E "\.8080\s" | grep -c ESTABLISHED)
    local TIME_WAIT=$(netstat -an 2>/dev/null | grep -E "\.8080\s" | grep -c TIME_WAIT)
    local TOTAL_TCP=$(netstat -an 2>/dev/null | grep -c tcp)
    
    # File descriptors
    local FDS=$(lsof -p $SERVER_PID 2>/dev/null | wc -l | tr -d ' ')
    
    # Update totals for averages
    TOTAL_CPU=$(echo "$TOTAL_CPU + ${CPU:-0}" | bc)
    TOTAL_RSS=$(echo "$TOTAL_RSS + ${RSS_MB:-0}" | bc)
    TOTAL_CONNS=$((TOTAL_CONNS + ESTABLISHED))
    SAMPLE_COUNT=$((SAMPLE_COUNT + 1))
    
    # Update peaks
    if (( $(echo "${CPU:-0} > $PEAK_CPU" | bc -l) )); then
        PEAK_CPU=${CPU:-0}
    fi
    if (( $(echo "${RSS_MB:-0} > $PEAK_RSS" | bc -l) )); then
        PEAK_RSS=${RSS_MB:-0}
    fi
    if [ "$ESTABLISHED" -gt "$PEAK_CONNS" ]; then
        PEAK_CONNS=$ESTABLISHED
    fi
    
    # Timestamp
    local TS=$(date +%H:%M:%S)
    
    # Print to console
    printf "\r[%s] CPU: %5.1f%% | RSS: %7.2f MB | Conns: %5d | FDs: %5d | Threads: %3d" \
        "$TS" "${CPU:-0}" "${RSS_MB:-0}" "$ESTABLISHED" "$FDS" "$THREADS"
    
    # Write to log
    {
        echo "[$TS] Elapsed: ${ELAPSED}s"
        echo "  CPU:         ${CPU:-0}%"
        echo "  RSS:         ${RSS_MB:-0} MB"
        echo "  VSZ:         ${VSZ_MB:-0} MB"
        echo "  Threads:     $THREADS"
        echo "  Established: $ESTABLISHED"
        echo "  TIME_WAIT:   $TIME_WAIT"
        echo "  FDs:         $FDS"
        echo "  Total TCP:   $TOTAL_TCP"
        echo ""
    } >> "$LOGFILE"
    
    # Write to CSV
    echo "$TS,$ELAPSED,${CPU:-0},${RSS_MB:-0},${VSZ_MB:-0},$THREADS,$ESTABLISHED,$TIME_WAIT,$FDS,$TOTAL_TCP" >> "$CSVFILE"
}

# Main capture loop
ELAPSED=0
while [ $ELAPSED -lt $DURATION ]; do
    # Check if process still exists
    if ! kill -0 $SERVER_PID 2>/dev/null; then
        echo ""
        echo "⚠️  Process $SERVER_PID terminated"
        break
    fi
    
    capture_sample $ELAPSED
    sleep $INTERVAL
    ELAPSED=$((ELAPSED + INTERVAL))
done

echo ""
echo ""

# ═══════════════════════════════════════════════════════════════════════════
# Calculate summary statistics
# ═══════════════════════════════════════════════════════════════════════════

if [ $SAMPLE_COUNT -gt 0 ]; then
    AVG_CPU=$(echo "scale=2; $TOTAL_CPU / $SAMPLE_COUNT" | bc)
    AVG_RSS=$(echo "scale=2; $TOTAL_RSS / $SAMPLE_COUNT" | bc)
    AVG_CONNS=$(echo "scale=0; $TOTAL_CONNS / $SAMPLE_COUNT" | bc)
    
    # Calculate per-user metrics (using peak connections as user count)
    if [ $PEAK_CONNS -gt 0 ]; then
        CPU_PER_USER=$(echo "scale=4; $PEAK_CPU / $PEAK_CONNS" | bc)
        RSS_PER_USER_KB=$(echo "scale=2; ($PEAK_RSS - $BASELINE_RSS_MB) * 1024 / $PEAK_CONNS" | bc)
        RSS_PER_USER_MB=$(echo "scale=4; ($PEAK_RSS - $BASELINE_RSS_MB) / $PEAK_CONNS" | bc)
    else
        CPU_PER_USER="N/A"
        RSS_PER_USER_KB="N/A"
        RSS_PER_USER_MB="N/A"
    fi
    
    # Get system info
    TOTAL_RAM_MB=$(sysctl -n hw.memsize 2>/dev/null | awk '{print $1/1024/1024}')
    CPU_CORES=$(sysctl -n hw.ncpu 2>/dev/null)
    FD_LIMIT=$(ulimit -n)
    SOMAXCONN=$(sysctl -n kern.ipc.somaxconn 2>/dev/null)
    
    # Calculate projected capacity
    if [ "$RSS_PER_USER_KB" != "N/A" ] && [ $(echo "$RSS_PER_USER_KB > 0" | bc) -eq 1 ]; then
        # RAM-based max (using 80% of total RAM)
        AVAILABLE_RAM=$(echo "scale=0; $TOTAL_RAM_MB * 0.8" | bc)
        MAX_USERS_RAM=$(echo "scale=0; ($AVAILABLE_RAM - $BASELINE_RSS_MB) * 1024 / $RSS_PER_USER_KB" | bc)
    else
        MAX_USERS_RAM="N/A"
    fi
    
    if [ "$CPU_PER_USER" != "N/A" ] && [ $(echo "$CPU_PER_USER > 0" | bc) -eq 1 ]; then
        # CPU-based max (using 80% of all cores)
        AVAILABLE_CPU=$(echo "scale=0; $CPU_CORES * 80" | bc)
        MAX_USERS_CPU=$(echo "scale=0; $AVAILABLE_CPU / $CPU_PER_USER" | bc)
    else
        MAX_USERS_CPU="N/A"
    fi
    
    # FD-based max
    MAX_USERS_FD=$((FD_LIMIT - 100))
    
    # Overall max is minimum of all limits
    # (simplified - just report each separately)
    
    # Generate report
    REPORT=$(cat << EOF

═══════════════════════════════════════════════════════════════════════════════
📊 SERVER CAPACITY ANALYSIS REPORT
═══════════════════════════════════════════════════════════════════════════════

SYSTEM INFORMATION
───────────────────────────────────────────────────────────────────────────────
  Total RAM:        ${TOTAL_RAM_MB} MB
  CPU Cores:        ${CPU_CORES}
  FD Limit:         ${FD_LIMIT}
  somaxconn:        ${SOMAXCONN}

BASELINE METRICS (before load)
───────────────────────────────────────────────────────────────────────────────
  Baseline RSS:     ${BASELINE_RSS_MB} MB

OBSERVED METRICS (during ${SAMPLE_COUNT} samples over ${DURATION}s)
───────────────────────────────────────────────────────────────────────────────
                    AVERAGE         PEAK
  CPU:              ${AVG_CPU}%          ${PEAK_CPU}%
  RSS:              ${AVG_RSS} MB       ${PEAK_RSS} MB
  Connections:      ${AVG_CONNS}            ${PEAK_CONNS}

PER-USER RESOURCE CONSUMPTION (at peak ${PEAK_CONNS} users)
───────────────────────────────────────────────────────────────────────────────
  CPU per user:     ${CPU_PER_USER}%
  RAM per user:     ${RSS_PER_USER_KB} KB (${RSS_PER_USER_MB} MB)

PROJECTED CAPACITY LIMITS
───────────────────────────────────────────────────────────────────────────────
  Max users (RAM @ 80%):    ${MAX_USERS_RAM}
  Max users (CPU @ 80%):    ${MAX_USERS_CPU}
  Max users (FD limit):     ${MAX_USERS_FD}
  Max users (somaxconn):    ${SOMAXCONN}

RECOMMENDED SERVER SPECS BY USER COUNT
───────────────────────────────────────────────────────────────────────────────
┌────────────┬────────────┬────────────┬────────────────────────────────────┐
│ Users      │ RAM Needed │ CPU Cores  │ Notes                              │
├────────────┼────────────┼────────────┼────────────────────────────────────┤
EOF
)

    # Calculate spec table dynamically
    for USERS in 1000 2500 5000 10000 25000 50000 100000; do
        if [ "$RSS_PER_USER_KB" != "N/A" ] && [ "$CPU_PER_USER" != "N/A" ]; then
            RAM_NEEDED=$(echo "scale=0; $BASELINE_RSS_MB + ($USERS * $RSS_PER_USER_KB / 1024)" | bc)
            CPU_NEEDED=$(echo "scale=1; $USERS * $CPU_PER_USER / 100" | bc)
            CPU_CORES_NEEDED=$(echo "scale=0; ($CPU_NEEDED + 0.9) / 1" | bc)  # Round up
            
            if [ $CPU_CORES_NEEDED -lt 1 ]; then
                CPU_CORES_NEEDED=1
            fi
            
            # Add notes based on scale
            if [ $USERS -le $PEAK_CONNS ]; then
                NOTES="✅ Tested stable"
            elif [ $USERS -le 5000 ]; then
                NOTES="Likely achievable"
            elif [ $USERS -le 20000 ]; then
                NOTES="Needs optimization"
            else
                NOTES="Consider clustering"
            fi
            
            REPORT="${REPORT}
│ $(printf '%10d' $USERS) │ $(printf '%7d MB' $RAM_NEEDED) │ $(printf '%10.0f' $CPU_CORES_NEEDED) │ $(printf '%-34s' "$NOTES") │"
        fi
    done

    REPORT="${REPORT}
└────────────┴────────────┴────────────┴────────────────────────────────────┘

FILES GENERATED
───────────────────────────────────────────────────────────────────────────────
  Detailed log:     $LOGFILE
  CSV data:         $CSVFILE

═══════════════════════════════════════════════════════════════════════════════
EOF

    # Print and save report
    echo "$REPORT"
    echo "$REPORT" >> "$LOGFILE"
    
    # Also save a summary file
    SUMMARY_FILE="capacity_summary_$(date +%Y%m%d_%H%M%S).txt"
    echo "$REPORT" > "$SUMMARY_FILE"
    echo ""
    echo "📁 Summary saved to: $SUMMARY_FILE"

else
    echo "❌ No samples captured"
fi