#!/bin/bash

# SYSTEM RESOURCE MONITOR - Agent 125
# Monitors system resources and prevents crashes during ML training

set -e

# Configuration
MEMORY_THRESHOLD=90          # Alert if memory usage >90%
SWAP_THRESHOLD=6144          # Alert if swap usage >6GB (in MB)
DISK_THRESHOLD=85            # Alert if disk usage >85%
CHECK_INTERVAL=60            # Check every 60 seconds
LOG_FILE="system_resource_monitor.log"
REPORT_FILE="SYSTEM_RESOURCE_MONITOR_REPORT.md"

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

# Alert counters
MEMORY_ALERTS=0
SWAP_ALERTS=0
DISK_ALERTS=0

# Function to log with timestamp
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

# Function to print colored output
print_status() {
    local color=$1
    local message=$2
    echo -e "${color}${message}${NC}"
    log "$message"
}

# Function to check memory usage
check_memory() {
    local mem_info=$(free | grep Mem)
    local total=$(echo $mem_info | awk '{print $2}')
    local used=$(echo $mem_info | awk '{print $3}')
    local free=$(echo $mem_info | awk '{print $4}')
    local available=$(echo $mem_info | awk '{print $7}')
    local percent=$((used * 100 / total))

    echo "$percent|$used|$total|$available"
}

# Function to check swap usage
check_swap() {
    local swap_info=$(free | grep Swap)
    local total=$(echo $swap_info | awk '{print $2}')
    local used=$(echo $swap_info | awk '{print $3}')

    if [ "$total" -eq 0 ]; then
        echo "0|0|0"
    else
        local percent=$((used * 100 / total))
        local used_mb=$((used / 1024))
        echo "$percent|$used_mb|$((total / 1024))"
    fi
}

# Function to check disk usage
check_disk() {
    local disk_info=$(df -h / | tail -1)
    local percent=$(echo $disk_info | awk '{print $5}' | tr -d '%')
    local used=$(echo $disk_info | awk '{print $3}')
    local total=$(echo $disk_info | awk '{print $2}')

    echo "$percent|$used|$total"
}

# Function to get training process info
get_training_processes() {
    # Look for ML training processes
    ps aux | grep -E "(train_liquid|train_|tune|optuna|gpu_training)" | grep -v grep || echo ""
}

# Function to get process memory usage
get_process_memory() {
    local pid=$1
    ps -p $pid -o rss= 2>/dev/null | awk '{printf "%.2f", $1/1024}' || echo "0"
}

# Function to get top memory consumers
get_top_memory_processes() {
    ps aux --sort=-%mem | head -11 | tail -10
}

# Function to generate alert recommendations
generate_recommendations() {
    local mem_percent=$1
    local swap_mb=$2
    local disk_percent=$3

    echo ""
    echo "## ALERT RECOMMENDATIONS"
    echo ""

    if [ "$mem_percent" -ge "$MEMORY_THRESHOLD" ]; then
        echo "### CRITICAL: High Memory Usage ($mem_percent%)"
        echo ""
        echo "**Immediate Actions:**"
        echo "1. Identify and kill non-essential processes"
        echo "2. Reduce batch size in training configuration"
        echo "3. Enable gradient checkpointing to reduce memory"
        echo "4. Consider using mixed precision (fp16) training"
        echo ""
        echo "**Process Kill Recommendations:**"
        echo '```bash'
        echo "# Kill non-essential Chrome/Firefox processes"
        echo "pkill -f 'chrome|firefox' 2>/dev/null || true"
        echo ""
        echo "# Kill Slack/Discord if running"
        echo "pkill -f 'slack|discord' 2>/dev/null || true"
        echo ""
        echo "# If still critical, reduce training batch size"
        echo "# Edit ml/examples/train_*.rs and reduce batch_size parameter"
        echo '```'
        echo ""
    fi

    if [ "$swap_mb" -ge "$SWAP_THRESHOLD" ]; then
        echo "### WARNING: High Swap Usage (${swap_mb}MB)"
        echo ""
        echo "**Immediate Actions:**"
        echo "1. System is thrashing - performance severely degraded"
        echo "2. Kill largest memory consumer immediately"
        echo "3. Restart training with reduced batch size"
        echo ""
        echo "**Emergency Kill Command:**"
        echo '```bash'
        echo "# Kill the largest memory consumer"
        echo "kill -9 \$(ps aux --sort=-%mem | awk 'NR==2 {print \$2}')"
        echo '```'
        echo ""
    fi

    if [ "$disk_percent" -ge "$DISK_THRESHOLD" ]; then
        echo "### WARNING: High Disk Usage ($disk_percent%)"
        echo ""
        echo "**Immediate Actions:**"
        echo "1. Clean up old checkpoints: rm -rf ml/tuning_checkpoints/trial_*/checkpoint_epoch_*"
        echo "2. Remove old logs: find . -name '*.log' -mtime +7 -delete"
        echo "3. Clean cargo cache: cargo clean"
        echo "4. Remove Docker volumes: docker system prune -a"
        echo ""
    fi
}

# Function to generate monitoring report
generate_report() {
    local duration=$1
    local max_mem=$2
    local max_swap=$3
    local max_disk=$4

    cat > "$REPORT_FILE" << EOF
# SYSTEM RESOURCE MONITOR REPORT - Agent 125

**Generated**: $(date '+%Y-%m-%d %H:%M:%S')
**Monitoring Duration**: ${duration} seconds
**Status**: $([ $MEMORY_ALERTS -eq 0 ] && [ $SWAP_ALERTS -eq 0 ] && [ $DISK_ALERTS -eq 0 ] && echo "✅ HEALTHY" || echo "⚠️ ALERTS DETECTED")

---

## Executive Summary

This report provides continuous system resource monitoring to prevent crashes during ML training.

### Alert Summary

| Resource | Alerts | Max Usage | Threshold | Status |
|----------|--------|-----------|-----------|--------|
| Memory   | $MEMORY_ALERTS | $max_mem% | $MEMORY_THRESHOLD% | $([ $MEMORY_ALERTS -eq 0 ] && echo "✅ OK" || echo "⚠️ CRITICAL") |
| Swap     | $SWAP_ALERTS | ${max_swap}MB | ${SWAP_THRESHOLD}MB | $([ $SWAP_ALERTS -eq 0 ] && echo "✅ OK" || echo "⚠️ WARNING") |
| Disk     | $DISK_ALERTS | $max_disk% | $DISK_THRESHOLD% | $([ $DISK_ALERTS -eq 0 ] && echo "✅ OK" || echo "⚠️ WARNING") |

---

## Current System Status

### Memory Status
\`\`\`
$(free -h)
\`\`\`

### Disk Status
\`\`\`
$(df -h /)
\`\`\`

### Top Memory Consumers
\`\`\`
$(get_top_memory_processes)
\`\`\`

### Active Training Processes
\`\`\`
$(get_training_processes || echo "No active training processes detected")
\`\`\`

---

## Resource Usage Timeline

**Note**: See $LOG_FILE for detailed timeline with timestamps

### Memory Usage Pattern
- Peak Memory: $max_mem%
- Alert Threshold: $MEMORY_THRESHOLD%
- Alerts Triggered: $MEMORY_ALERTS

### Swap Usage Pattern
- Peak Swap: ${max_swap}MB
- Alert Threshold: ${SWAP_THRESHOLD}MB
- Alerts Triggered: $SWAP_ALERTS

### Disk Usage Pattern
- Peak Disk: $max_disk%
- Alert Threshold: $DISK_THRESHOLD%
- Alerts Triggered: $DISK_ALERTS

---

## Monitoring Configuration

\`\`\`yaml
check_interval: ${CHECK_INTERVAL}s
thresholds:
  memory: ${MEMORY_THRESHOLD}%
  swap: ${SWAP_THRESHOLD}MB
  disk: ${DISK_THRESHOLD}%
alerts:
  memory: $([ $MEMORY_ALERTS -eq 0 ] && echo "none" || echo "$MEMORY_ALERTS triggered")
  swap: $([ $SWAP_ALERTS -eq 0 ] && echo "none" || echo "$SWAP_ALERTS triggered")
  disk: $([ $DISK_ALERTS -eq 0 ] && echo "none" || echo "$DISK_ALERTS triggered")
\`\`\`

---

## Recommendations

### System Optimization

1. **Memory Management**:
   - Current available: $(free -h | grep Mem | awk '{print $7}')
   - Recommendation: $([ "$max_mem" -lt 70 ] && echo "✅ Healthy - no action needed" || echo "⚠️ Consider reducing batch size or enabling gradient checkpointing")

2. **Swap Usage**:
   - Current swap: $(free -h | grep Swap | awk '{print $3}')
   - Recommendation: $([ "$max_swap" -lt 1024 ] && echo "✅ Minimal swap - good performance" || echo "⚠️ High swap indicates memory pressure - upgrade RAM or reduce workload")

3. **Disk Space**:
   - Available: $(df -h / | tail -1 | awk '{print $4}')
   - Recommendation: $([ "$max_disk" -lt 70 ] && echo "✅ Sufficient space" || echo "⚠️ Clean up old checkpoints and logs")

### Training Optimization

1. **Batch Size Tuning**:
   - If memory >80%, reduce batch_size by 50%
   - If memory <50%, can increase batch_size by 50%

2. **Gradient Checkpointing**:
   - Enable for MAMBA-2/TFT models if memory >70%
   - Trades 30% more compute for 50% less memory

3. **Mixed Precision**:
   - Use fp16 instead of fp32 to halve memory usage
   - Minimal accuracy impact (<0.5% for most models)

---

## Emergency Procedures

### If Memory >90%

\`\`\`bash
# 1. Kill non-essential processes
pkill -f 'chrome|firefox|slack' 2>/dev/null || true

# 2. Clear page cache (safe, no data loss)
sudo sync && sudo sh -c 'echo 3 > /proc/sys/vm/drop_caches'

# 3. Kill largest memory consumer (emergency only)
kill -9 \$(ps aux --sort=-%mem | awk 'NR==2 {print \$2}')
\`\`\`

### If Swap >6GB

\`\`\`bash
# System is thrashing - immediate action required
# Kill the training process and restart with reduced batch size
pkill -f 'train_liquid|optuna'

# Wait for swap to clear
sleep 30

# Reduce batch size in config and restart
\`\`\`

### If Disk >85%

\`\`\`bash
# Clean up old checkpoints
rm -rf ml/tuning_checkpoints/trial_*/checkpoint_epoch_*

# Remove old logs
find . -name '*.log' -mtime +7 -delete

# Clean cargo cache
cargo clean

# Check space again
df -h /
\`\`\`

---

## Continuous Monitoring

**Status**: $([ -f "/tmp/resource_monitor.pid" ] && echo "🟢 RUNNING" || echo "🔴 STOPPED")

To continue monitoring:
\`\`\`bash
# Start monitoring (runs in background)
nohup ./scripts/system_resource_monitor.sh monitor > /dev/null 2>&1 &

# Stop monitoring
kill \$(cat /tmp/resource_monitor.pid 2>/dev/null)
\`\`\`

To generate this report again:
\`\`\`bash
./scripts/system_resource_monitor.sh report
\`\`\`

---

## Log File

Full monitoring log: \`$LOG_FILE\`

\`\`\`bash
# View last 50 entries
tail -50 $LOG_FILE

# View all alerts
grep -E 'ALERT|WARNING|CRITICAL' $LOG_FILE

# View memory timeline
grep 'Memory:' $LOG_FILE
\`\`\`

---

**Last Updated**: $(date '+%Y-%m-%d %H:%M:%S')
**Next Check**: In ${CHECK_INTERVAL} seconds (if monitoring active)
**Generated by**: Agent 125 - System Resource Monitor
EOF

    print_status "$GREEN" "✅ Report generated: $REPORT_FILE"
}

# Function to monitor resources continuously
monitor_resources() {
    local start_time=$(date +%s)
    local max_mem=0
    local max_swap=0
    local max_disk=0

    # Save PID for stopping
    echo $$ > /tmp/resource_monitor.pid

    print_status "$BLUE" "🚀 Starting system resource monitoring..."
    print_status "$BLUE" "⚙️  Check interval: ${CHECK_INTERVAL}s"
    print_status "$BLUE" "⚠️  Thresholds: Memory ${MEMORY_THRESHOLD}%, Swap ${SWAP_THRESHOLD}MB, Disk ${DISK_THRESHOLD}%"
    echo ""

    # Main monitoring loop
    while true; do
        # Check memory
        mem_data=$(check_memory)
        mem_percent=$(echo $mem_data | cut -d'|' -f1)
        mem_used=$(echo $mem_data | cut -d'|' -f2)
        mem_total=$(echo $mem_data | cut -d'|' -f3)
        mem_available=$(echo $mem_data | cut -d'|' -f4)

        [ $mem_percent -gt $max_mem ] && max_mem=$mem_percent

        # Check swap
        swap_data=$(check_swap)
        swap_percent=$(echo $swap_data | cut -d'|' -f1)
        swap_mb=$(echo $swap_data | cut -d'|' -f2)
        swap_total=$(echo $swap_data | cut -d'|' -f3)

        [ $swap_mb -gt $max_swap ] && max_swap=$swap_mb

        # Check disk
        disk_data=$(check_disk)
        disk_percent=$(echo $disk_data | cut -d'|' -f1)
        disk_used=$(echo $disk_data | cut -d'|' -f2)
        disk_total=$(echo $disk_data | cut -d'|' -f3)

        [ $disk_percent -gt $max_disk ] && max_disk=$disk_percent

        # Determine status color
        if [ $mem_percent -ge $MEMORY_THRESHOLD ] || [ $swap_mb -ge $SWAP_THRESHOLD ]; then
            STATUS_COLOR=$RED
            STATUS="CRITICAL"
        elif [ $mem_percent -ge 75 ] || [ $swap_mb -ge 4096 ] || [ $disk_percent -ge $DISK_THRESHOLD ]; then
            STATUS_COLOR=$YELLOW
            STATUS="WARNING"
        else
            STATUS_COLOR=$GREEN
            STATUS="OK"
        fi

        # Print status
        echo -e "${STATUS_COLOR}[$(date '+%H:%M:%S')] Memory: ${mem_percent}% | Swap: ${swap_mb}MB | Disk: ${disk_percent}% | Status: $STATUS${NC}"
        log "Memory: ${mem_percent}% (${mem_used}KB/${mem_total}KB, Available: ${mem_available}KB) | Swap: ${swap_mb}MB/${swap_total}MB | Disk: ${disk_percent}% (${disk_used}/${disk_total})"

        # Check for alerts
        if [ $mem_percent -ge $MEMORY_THRESHOLD ]; then
            ((MEMORY_ALERTS++))
            print_status "$RED" "🚨 MEMORY ALERT: ${mem_percent}% usage exceeds ${MEMORY_THRESHOLD}% threshold!"
            print_status "$RED" "   Available: $((mem_available / 1024))MB"

            # Show top processes
            print_status "$YELLOW" "   Top 5 memory consumers:"
            ps aux --sort=-%mem | head -6 | tail -5 | awk '{printf "   - %s (PID %s): %.1f%%\n", $11, $2, $4}' | tee -a "$LOG_FILE"

            # Generate recommendations
            generate_recommendations $mem_percent $swap_mb $disk_percent | tee -a "$LOG_FILE"
        fi

        if [ $swap_mb -ge $SWAP_THRESHOLD ]; then
            ((SWAP_ALERTS++))
            print_status "$RED" "🚨 SWAP ALERT: ${swap_mb}MB usage exceeds ${SWAP_THRESHOLD}MB threshold!"
            print_status "$RED" "   System is likely thrashing - performance severely degraded"
            generate_recommendations $mem_percent $swap_mb $disk_percent | tee -a "$LOG_FILE"
        fi

        if [ $disk_percent -ge $DISK_THRESHOLD ]; then
            ((DISK_ALERTS++))
            print_status "$YELLOW" "⚠️  DISK ALERT: ${disk_percent}% usage exceeds ${DISK_THRESHOLD}% threshold!"
            print_status "$YELLOW" "   Available: $(df -h / | tail -1 | awk '{print $4}')"
            generate_recommendations $mem_percent $swap_mb $disk_percent | tee -a "$LOG_FILE"
        fi

        # Check for training processes
        training_procs=$(get_training_processes)
        if [ -n "$training_procs" ]; then
            log "Active training processes detected:"
            echo "$training_procs" | while read -r line; do
                log "  $line"
            done
        fi

        # Generate report every 10 checks (10 minutes)
        if [ $(($(date +%s) - start_time)) -ge 600 ]; then
            generate_report $(($(date +%s) - start_time)) $max_mem $max_swap $max_disk
            start_time=$(date +%s)
        fi

        # Sleep until next check
        sleep $CHECK_INTERVAL
    done
}

# Function to show current status (one-time check)
show_status() {
    print_status "$BLUE" "📊 System Resource Status"
    echo ""

    # Memory
    mem_data=$(check_memory)
    mem_percent=$(echo $mem_data | cut -d'|' -f1)
    mem_used=$(echo $mem_data | cut -d'|' -f2)
    mem_total=$(echo $mem_data | cut -d'|' -f3)
    mem_available=$(echo $mem_data | cut -d'|' -f4)

    [ $mem_percent -ge $MEMORY_THRESHOLD ] && COLOR=$RED || [ $mem_percent -ge 75 ] && COLOR=$YELLOW || COLOR=$GREEN
    print_status "$COLOR" "Memory: ${mem_percent}% (Available: $((mem_available / 1024))MB)"

    # Swap
    swap_data=$(check_swap)
    swap_mb=$(echo $swap_data | cut -d'|' -f2)

    [ $swap_mb -ge $SWAP_THRESHOLD ] && COLOR=$RED || [ $swap_mb -ge 4096 ] && COLOR=$YELLOW || COLOR=$GREEN
    print_status "$COLOR" "Swap: ${swap_mb}MB"

    # Disk
    disk_data=$(check_disk)
    disk_percent=$(echo $disk_data | cut -d'|' -f1)
    disk_used=$(echo $disk_data | cut -d'|' -f2)
    disk_total=$(echo $disk_data | cut -d'|' -f3)

    [ $disk_percent -ge $DISK_THRESHOLD ] && COLOR=$RED || [ $disk_percent -ge 70 ] && COLOR=$YELLOW || COLOR=$GREEN
    print_status "$COLOR" "Disk: ${disk_percent}% (${disk_used}/${disk_total})"

    echo ""
    print_status "$BLUE" "Top 5 Memory Consumers:"
    ps aux --sort=-%mem | head -6 | tail -5 | awk '{printf "%s (PID %s): %.1f%% (%.1fMB)\n", $11, $2, $4, $6/1024}'

    echo ""
    print_status "$BLUE" "Active Training Processes:"
    training_procs=$(get_training_processes)
    if [ -n "$training_procs" ]; then
        echo "$training_procs"
    else
        echo "No active training processes detected"
    fi

    # Generate report
    generate_report 0 $mem_percent $swap_mb $disk_percent
}

# Main script logic
case "${1:-monitor}" in
    monitor)
        monitor_resources
        ;;
    status)
        show_status
        ;;
    report)
        show_status
        ;;
    stop)
        if [ -f /tmp/resource_monitor.pid ]; then
            pid=$(cat /tmp/resource_monitor.pid)
            kill $pid 2>/dev/null && print_status "$GREEN" "✅ Monitoring stopped" || print_status "$YELLOW" "⚠️  No monitoring process found"
            rm -f /tmp/resource_monitor.pid
        else
            print_status "$YELLOW" "⚠️  No monitoring process found"
        fi
        ;;
    *)
        echo "Usage: $0 {monitor|status|report|stop}"
        echo ""
        echo "Commands:"
        echo "  monitor  - Start continuous monitoring (default)"
        echo "  status   - Show current status (one-time check)"
        echo "  report   - Generate monitoring report"
        echo "  stop     - Stop monitoring process"
        exit 1
        ;;
esac