#!/bin/bash

# Script to monitor and log space, memory, and CPU usage of a process
# Usage: sudo ./monitor_process.sh <process_name> <log_file>

# Enable error handling after argument check
set -o errexit
set -o nounset
set -o pipefail

# Function to log message to both console and file
log() {
  echo "$1" | tee -a "$LOG_FILE"
}

# Function to show usage
usage() {
  echo "Usage: $0 <process_name> [log_file]"
  echo "  process_name: Name of the process to monitor"
  echo "  log_file: Optional. Path to log file (default: ./<process_name>_monitor.log)"
  exit 1
}

# Check for required arguments
if [ $# -lt 1 ]; then
  usage
fi

# Configuration
readonly PROCESS_NAME="$1"
readonly LOG_FILE="${2:-"./${PROCESS_NAME}_monitor.log"}"
readonly INTERVAL=5       # Monitoring interval in seconds
readonly MAX_RUNTIME=1800 # Maximum runtime in seconds (30 minutes)
readonly START_TIME=$(date +%s)
DISK_MEASUREMENTS=() # Array to store disk usage measurements

# Function to get process PID
get_pid() {
  pgrep -o -f "$PROCESS_NAME" || true
}

# Function to find process read/write locations
find_process_locations() {
  local pid=$1
  local req_mode="${2:-w}" # Optional: "r" for read-only, "w" for write (default), "rw" for read-write
  local fd_dir="/proc/$pid/fd"

  # Check if PID is provided
  if [ -z "$pid" ]; then
    echo "Error: PID not provided" >&2
    return 1
  fi

  # Check if process exists
  if ! kill -0 "$pid" 2>/dev/null; then
    echo "Error: Process $pid does not exist" >&2
    return 1
  fi

  # Check if fd directory exists and is readable
  if [ ! -d "$fd_dir" ] || [ ! -r "$fd_dir" ]; then
    echo "Error: Cannot access $fd_dir" >&2
    return 1
  fi

  # Find and filter file descriptors opened for writing
  find "$fd_dir" -type l 2>/dev/null | while read -r fd; do
    local fdnum
    fdnum=$(basename "$fd")
    local info_file="/proc/$pid/fdinfo/$fdnum"
    # Check if fdinfo file exists
    [ -f "$info_file" ] || continue

    local flag
    flag=$(grep '^flags:' "$info_file" | awk '{print $2}')
    # Compute access mode: 0 = read-only, 1 = write-only, 2 = read-write
    local fdmode=$((flag & 3))

    # Check against desired mode
    case "$req_mode" in
    r)
      [ "$fdmode" -eq 0 ] || continue
      ;;
    w)
      [ "$fdmode" -ne 0 ] || continue
      ;;
    rw)
      [ "$fdmode" -eq 2 ] || continue
      ;;
    *)
      # If an unsupported mode is provided, skip this FD
      continue
      ;;
    esac

    local target
    target=$(readlink "$fd" 2>/dev/null) || continue

    # Skip special files and directories
    case "$target" in
    /dev/* | pipe:* | socket:* | anon_inode:*) continue ;;
    esac

    # Only show regular files and directories that exist
    if [ -f "$target" ] || [ -d "$target" ]; then
      echo "$target"
    fi
  done | sort -u
}


# Function to get disk usage
get_disk_usage() {
  local pid=$1
  if [ -d "/proc/$pid/cwd" ]; then
    du -sh "/proc/$pid/cwd" 2>/dev/null | awk '{print $1}' || echo "N/A"
  else
    echo "N/A"
  fi
}

# Function to get memory usage
get_memory_usage() {
  local pid=$1
  pmap "$pid" 2>/dev/null | tail -n 1 | awk '/[0-9]K/{print $2}' || echo "N/A"
}

# Function to get CPU usage
get_cpu_usage() {
  local pid=$1
  ps -p "$pid" -o %cpu= 2>/dev/null | awk '{printf "%.1f", $1}' || echo "N/A"
}

# Function to calculate average disk usage rate
calculate_average_disk_rate() {
  local -a measurements=("${@}")
  local count=${#measurements[@]}
  local total=0

  for ((i = 1; i < count; i++)); do
    local diff=$((measurements[i] - measurements[i - 1]))
    total=$((total + diff))
  done

  # Calculate average change per interval
  if [ $((count - 1)) -gt 0 ]; then
    echo $((total / (count - 1)))
  else
    echo 0
  fi
}

# Initialize log file with headers
{
  log "=== Process Monitoring Report ==="
  log "Started at: $(date)"
  log "Process: $PROCESS_NAME"
  log "Monitoring interval: ${INTERVAL}s"
  log "Maximum runtime: ${MAX_RUNTIME}s"
  log "----------------------------------------"
}

# Get initial measurements
PID=$(get_pid)
if [ -z "$PID" ]; then
  log "Error: Process $PROCESS_NAME not found"
  exit 1
fi

# Store initial disk usage
INITIAL_DISK_USAGE=$(get_disk_usage "$PID")
DISK_MEASUREMENTS+=("$(echo "$INITIAL_DISK_USAGE" | numfmt --from=iec)")
log "Initial disk usage: $INITIAL_DISK_USAGE"

# Monitoring loop
while true; do
  CURRENT_TIME=$(date +%s)
  RUNTIME=$((CURRENT_TIME - START_TIME))

  # Exit if maximum runtime is reached
  if [ "$RUNTIME" -ge "$MAX_RUNTIME" ]; then
    FINAL_DISK_USAGE=$(get_disk_usage "$PID")
    DISK_MEASUREMENTS+=("$(echo "$FINAL_DISK_USAGE" | numfmt --from=iec)")
    USAGE_RATE=$(calculate_disk_usage_rate "$INITIAL_DISK_USAGE" "$FINAL_DISK_USAGE" "$RUNTIME")
    AVG_RATE=$(calculate_average_disk_rate "${DISK_MEASUREMENTS[@]}")

    {
      log "=== Monitoring Summary ==="
      log "Ended at: $(date)"
      log "Total runtime: $RUNTIME seconds"
      log "Initial disk usage: $INITIAL_DISK_USAGE"
      log "Final disk usage: $FINAL_DISK_USAGE"
      log "Disk usage rate: ${USAGE_RATE} MB/minute"
      log "Average rate of change: ${AVG_RATE} bytes/interval"
    }
    exit 0
  fi

  # Get current PID
  PID=$(get_pid)
  if [ -z "$PID" ]; then
    log "$(date): Process $PROCESS_NAME not found"
    sleep "$INTERVAL"
    continue
  fi

  # Get current measurements and store disk usage
  MEM_USAGE=$(get_memory_usage "$PID")
  CPU_USAGE=$(get_cpu_usage "$PID")
  DISK_USAGE=$(get_disk_usage "$PID")
  DISK_MEASUREMENTS+=("$(echo "$DISK_USAGE" | numfmt --from=iec)")

  # Log the information
  log "$(date): PID=$PID, Memory=$MEM_USAGE, CPU=${CPU_USAGE}%, Disk=$DISK_USAGE"

  sleep "$INTERVAL"
done