#!/bin/bash
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Exit on error, treat unset variables as errors, and propagate pipeline errors.
set -euo pipefail

# This script is used to collect the resource(mem,cpu,disk) usage of the machine in the file at given interval.
# This script can also be in use to print the resource usage visualisation in graph format.
usage() {
    echo "A script to collect and visualize system resource usage."
    echo ""
    echo "========================================================================"
    echo ""
    echo "COMMAND: COLLECT"
    echo ""
    echo "  Gathers resource usage (CPU, memory, disk) at predefined intervals."
    echo ""
    echo "  USAGE:"
    echo "    $0 COLLECT <FILE_PATH>"
    echo ""
    echo "  ARGUMENTS:"
    echo ""
    echo "    <FILE_PATH>   The path to the output file where resource data will be"
    echo "                  saved. The script will create or overwrite this file."
    echo ""
    echo "========================================================================"
    echo ""
    echo "COMMAND: PRINT"
    echo ""
    echo "  Visualizes resource data from a file in a graph format."
    echo ""
    echo "  USAGE:"
    echo "    $0 PRINT <FILE_PATH>"
    echo ""
    echo "  ARGUMENTS:"
    echo "    <FILE_PATH>   The path to the input data file (generated by COLLECT)"
    echo "                  that contains the statistics to be visualized."
    echo ""
    echo "========================================================================"
    exit 1
}

readonly USAGE_COLLECTION_INTERVAL_SECONDS=10               # Collect usage every 10s.
readonly USAGE_COLLECTION_DURATION_SECONDS=$((3 * 60 * 60)) # Auto collect usage for upto 3 hrs only (VM timeout is 3 hrs in kokoro).
readonly COLLECT="COLLECT"
readonly PRINT="PRINT"

# Validate number of arguments.
if [[ $# -ne 2 ]]; then
    echo "Error: Invalid number of arguments."
    usage
fi

readonly COMMAND="$1"
readonly FILE_PATH="$2"

# Calculates CPU usage as integer percentage using top command.
get_cpu_usage() {
    # 'top -b -n 1 | awk '/^%Cpu/'' returns
    # %Cpu(s):  1.1 us,  0.0 sy,  0.0 ni, 98.9 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
    # 8th entry is idle cpu %
    top -b -n 1 | awk '/^%Cpu/ { usage = 100 - $8; print int(usage) }' || echo 0
}

# Calculates memory usage as integer percentage using free.
get_mem_usage() {
    # 'free' returns
    #                total        used        free      shared  buff/cache   available
    # Mem:       371458720     6879376   285303364        9772    81770240   364579344
    # Swap:              0           0           0
    free | awk '/Mem:/ {p=($2-$7)*100/$2; print int(p) }' || echo 0
}

# Calculates disk usage as integer percentage using df.
get_disk_usage() {
    # 'df -P /' returns
    # Filesystem     1024-blocks     Used Available Capacity Mounted on
    # /dev/nvme0n1p1   515794480 82928236 411780644      17% /
    df -P / | tail -n1 | awk '{sub(/%/, ""); printf "%.0f\n", $5}' || echo 0
}

# Helper function to collect the usage and write to file.
# Output: Appends new line to file with %cpu %mem %disk at given interval.
collect_all_metric() {
    echo "CPU MEM DISK" >"$FILE_PATH"
    while ((${SECONDS:-0} < USAGE_COLLECTION_DURATION_SECONDS)); do
        {
            cpu=$(get_cpu_usage)
            mem=$(get_mem_usage)
            disk=$(get_disk_usage)
            echo "$cpu $mem $disk"
        } >>"$FILE_PATH"
        sleep "$USAGE_COLLECTION_INTERVAL_SECONDS"
    done
    echo "Exiting due to monitoring time limit."
}

# Helper function to print usage of all metrics.
print_each_metric() {
    # Read headers from file.
    read -r -a headers <"$FILE_PATH"
    declare -A data_columns
    local line_num=0
    while read -r -a values; do
        # Skip the header line we already read
        ((line_num++ == 0)) && continue

        for i in "${!headers[@]}"; do
            local header="${headers[$i]}"
            local value="${values[$i]}"
            # Append the value to the correct string in the associative array
            data_columns["$header"]+="$value "
        done
    done <"$FILE_PATH"

    # Iterate through each resource and print its graph
    for resource in "${headers[@]}"; do
        local -a usage_values=(${data_columns[$resource]})
        local max_value=0
        for val in "${usage_values[@]}"; do
            ((val > max_value)) && max_value=$val
        done

        # Add some buffer to the graph ceiling, capping at 100
        local graph_ceiling=$((max_value + 5))
        ((graph_ceiling > 100)) && graph_ceiling=100

        printf "\n%s usage in percentage every %s seconds:\n\n" "$resource" "$USAGE_COLLECTION_INTERVAL_SECONDS"

        # Print the graph from top to bottom
        for ((y = graph_ceiling; y >= 0; y--)); do
            printf "%s@ %3d%%|" "$resource" "$y"
            for val in "${usage_values[@]}"; do
                if ((val >= y)); then
                    printf "*"
                else
                    printf " "
                fi
            done
            printf "\n"
        done

        # Print separator line
        local total_width=$((${#usage_values[@]} + 15))
        for ((i = 0; i < total_width; i++)); do
            printf "="
        done
        printf "\n\n"
    done
}

# Exit the script if SIGTERM or SIGINT is received.
trap 'exit 0' SIGINT SIGTERM

if [[ "$COMMAND" == "$COLLECT" ]]; then
    collect_all_metric
elif [[ "$COMMAND" == "$PRINT" ]]; then
    print_each_metric
else
    echo "Error: Invalid command."
    usage
fi
