#!/usr/bin/env bash

set -u

KEYWORD="${1:-}"
KEYWORD_LOWER="${KEYWORD,,}"

if [[ -z "$KEYWORD" ]]; then
    echo "Usage: $0 <keyword>"
    echo
    echo "Examples:"
    echo "  $0 firefox"
    echo "  $0 chrome"
    exit 1
fi

TOTAL_CPU=0
TOTAL_RAM_KB=0
TOTAL_SWAP_KB=0
PROCESS_COUNT=0

# ps aux:
#
# USER PID %CPU %MEM VSZ RSS TTY STAT START TIME COMMAND
#
# We skip the header and inspect:
#
# $2 = PID
# $3 = CPU
# $6 = RSS
# $11 onward = COMMAND
#
# Using "read -ra" would lose command boundaries, so we use awk
# to print the first 10 fields plus the complete command.

while read -r pid cpu rss command; do

    [[ -n "$pid" ]] || continue

    # Case-insensitive matching against the complete command.
    if [[ "${command,,}" != *"$KEYWORD_LOWER"* ]]; then
        continue
    fi

    PROCESS_COUNT=$((PROCESS_COUNT + 1))

    # CPU percentage from ps.
    TOTAL_CPU=$(awk \
        -v a="$TOTAL_CPU" \
        -v b="$cpu" \
        'BEGIN { printf "%.2f", a + b }')

    # RSS is in KB.
    TOTAL_RAM_KB=$((TOTAL_RAM_KB + rss))

    # VmSwap is in kB.
    if [[ -r "/proc/$pid/status" ]]; then
        swap=$(awk '
            /^VmSwap:/ {
                print $2
                found=1
                exit
            }
            END {
                if (!found)
                    print 0
            }
        ' "/proc/$pid/status" 2>/dev/null)

        [[ "$swap" =~ ^[0-9]+$ ]] || swap=0

        TOTAL_SWAP_KB=$((TOTAL_SWAP_KB + swap))
    fi

done < <(
    ps aux |
    tail -n +2 |
    awk '
    {
        pid=$2
        cpu=$3
        rss=$6

        command=""
        for (i=11; i<=NF; i++) {
            command = command (i==11 ? "" : " ") $i
        }

        print pid, cpu, rss, command
    }'
)

# ------------------------------------------------------------
# Format KB
# ------------------------------------------------------------

format_kb() {
    local kb="$1"

    if (( kb >= 1024 * 1024 )); then
        awk -v kb="$kb" \
            'BEGIN { printf "%.2f GB", kb / 1024 / 1024 }'

    elif (( kb >= 1024 )); then
        awk -v kb="$kb" \
            'BEGIN { printf "%.0f MB", kb / 1024 }'

    else
        printf "%d KB" "$kb"
    fi
}

echo
echo "$KEYWORD"
echo "  Processes : $PROCESS_COUNT"
printf "  CPU       : %.2f%%\n" "$TOTAL_CPU"
echo "  RAM       : $(format_kb "$TOTAL_RAM_KB")"
echo "  Swap      : $(format_kb "$TOTAL_SWAP_KB")"
echo