#!/bin/bash

# This script monitors the total memory usage of all Chrome processes
# and displays an ASCII line chart of the usage over the last 2 hours,
# ensuring the chart always fits within the terminal width.

# Configuration
INTERVAL_SECONDS=5 # How often to sample memory usage (in seconds)
DURATION_HOURS=2    # How many hours of data to display on the chart
DATA_FILE="/tmp/chrome_memory_data.txt" # Temporary file to store historical data
CHART_HEIGHT=20     # Number of lines for the chart's vertical height (excluding labels)

# --- Script Initialization ---

# Clear the terminal when the script starts
tput clear

# Ensure the temporary data file exists. If it doesn't, touch creates it.
touch "$DATA_FILE"

# --- Functions ---

# Function to get the total memory usage of all Chrome processes in MB.
# It sums the Resident Set Size (RSS) for all processes named 'chrome'.
get_chrome_memory_usage() {
    # Find all process IDs (PIDs) for processes named 'chrome'.
    # '-d,' makes pgrep output PIDs separated by commas, suitable for 'ps -p'.
    local pids=$(pgrep -d, chrome)

    # If no Chrome processes are found, return 0.
    if [ -z "$pids" ]; then
        echo "0"
        return
    fi

    # Sum the RSS (Resident Set Size) in KB for all identified Chrome PIDs.
    # 'ps -p "$pids" -o rss=' outputs only the RSS value in KB for each PID.
    # '2>/dev/null' suppresses errors if some PIDs no longer exist.
    # 'awk' sums all the RSS values.
    local total_rss_kb=$(ps -p "$pids" -o rss= 2>/dev/null | awk '{s+=$1} END {print s}')

    # If no memory is reported (e.g., processes died between pgrep and ps), return 0.
    if [ -z "$total_rss_kb" ]; then
        echo "0"
        return
    fi

    # Convert KB to MB (1024 KB = 1 MB) using 'bc' for floating-point arithmetic.
    local total_rss_mb=$(echo "scale=2; $total_rss_kb / 1024" | bc)
    echo "$total_rss_mb"
}

# Function to clean old data from the temporary file.
# It removes any entries older than the specified DURATION_HOURS
# and also ensures the number of data points does not exceed terminal width.
clean_old_data() {
    local current_timestamp=$(date +%s) # Get current Unix timestamp
    # Calculate the cutoff timestamp: current time minus the duration in seconds.
    local cutoff_timestamp=$((current_timestamp - DURATION_HOURS * 3600)) # 3600 seconds in an hour

    # First, filter data based on time.
    awk -F':' -v cutoff="$cutoff_timestamp" '$1 >= cutoff' "$DATA_FILE" > "${DATA_FILE}.tmp" && \
    mv "${DATA_FILE}.tmp" "$DATA_FILE"

    # Now, ensure the number of data points does not exceed terminal width.
    # Get current terminal columns.
    local current_columns=$(tput cols)
    # Estimate the width taken by the Y-axis label and padding.
    # "  999.9 MB |" is roughly 14 characters.
    local y_axis_label_width=14
    # Calculate maximum available columns for the chart data points.
    local max_chart_columns=$((current_columns - y_axis_label_width - 2)) # -2 for extra padding/border

    # Ensure a minimum number of columns to avoid negative values or very small charts.
    if [ "$max_chart_columns" -lt 10 ]; then
        max_chart_columns=10 # Ensure at least 10 columns for readability
    fi

    # Read the data again after time-based cleaning.
    local filtered_data=()
    while IFS=':' read -r timestamp mem_mb; do
        filtered_data+=("$timestamp:$mem_mb")
    done < "$DATA_FILE"

    local num_filtered_points=${#filtered_data[@]}

    # If the number of data points exceeds the maximum chart columns,
    # keep only the latest 'max_chart_columns' data points.
    if [ "$num_filtered_points" -gt "$max_chart_columns" ]; then
        local start_index=$((num_filtered_points - max_chart_columns))
        printf "%s\n" "${filtered_data[@]:$start_index:$max_chart_columns}" > "$DATA_FILE"
    fi
}

# Function to draw the ASCII line chart in the terminal.
draw_chart() {
    local data=()          # Array to store memory usage values
    local timestamps=()    # Array to store corresponding timestamps
    local max_mem=0        # Maximum memory usage found in the current data
    local min_mem=9999999  # Minimum memory usage found (initialized high)

    # Read data from the file into the arrays and find min/max memory.
    while IFS=':' read -r timestamp mem_mb; do
        timestamps+=("$timestamp")
        data+=("$mem_mb")
        # Update max_mem if current value is higher.
        if (( $(echo "$mem_mb > $max_mem" | bc -l) )); then
            max_mem="$mem_mb"
        fi
        # Update min_mem if current value is lower.
        if (( $(echo "$mem_mb < $min_mem" | bc -l) )); then
            min_mem="$mem_mb"
        fi
    done < "$DATA_FILE"

    local num_data_points=${#data[@]} # Total number of data points to chart

    # If no data is available, print a message and exit the function.
    if [ "$num_data_points" -eq 0 ]; then
        echo "No data to display yet. Waiting for Chrome memory usage..."
        return
    fi

    # Adjust max_mem and min_mem for better chart scaling and padding.
    # Ensure a minimum range if memory usage is very low to prevent flat charts.
    if (( $(echo "$max_mem < 10" | bc -l) )); then
        max_mem=10
    fi
    # If min_mem is very close to max_mem, expand the range slightly for visibility.
    if (( $(echo "$max_mem - $min_mem < 5" | bc -l) )); then
        local avg_mem=$(echo "scale=2; ($max_mem + $min_mem) / 2" | bc)
        max_mem=$(echo "scale=2; $avg_mem + 2.5" | bc)
        min_mem=$(echo "scale=2; $avg_mem - 2.5" | bc)
        # Ensure min_mem doesn't go below zero.
        if (( $(echo "$min_mem < 0" | bc -l) )); then min_mem=0; fi
    fi

    # Calculate the total memory range for scaling.
    local mem_range=$(echo "scale=2; $max_mem - $min_mem" | bc)
    # Prevent division by zero if all memory values are identical.
    if (( $(echo "$mem_range == 0" | bc -l) )); then
        mem_range=1
    fi

    echo "Chrome Memory Usage (last ${DURATION_HOURS} hours)"
    echo "--------------------------------------------------"

    # Array to store the scaled vertical position (row index) for each data point.
    local scaled_positions=()
    for (( i = 0; i < num_data_points; i++ )); do
        local current_mem="${data[i]}"
        # Scale the current memory value to a position within the CHART_HEIGHT (0 to CHART_HEIGHT).
        local scaled_pos=$(echo "scale=0; ($current_mem - $min_mem) * $CHART_HEIGHT / $mem_range" | bc)
        scaled_positions+=("$scaled_pos")
    done

    # Draw the chart rows from top to bottom (highest memory level to lowest).
    for (( h = CHART_HEIGHT; h >= 0; h-- )); do
        local line=""
        # Calculate the memory level corresponding to the current row 'h'.
        local mem_level_label=$(echo "scale=1; $min_mem + ($mem_range * $h / $CHART_HEIGHT)" | bc)
        # Print the Y-axis label (memory level).
        LC_NUMERIC="C" printf "%6.1f MB |" "$mem_level_label"

        # Iterate through each data point to draw the chart line.
        for (( i = 0; i < num_data_points; i++ )); do
            # If the scaled position of the data point matches the current row 'h', print an asterisk.
            if [ "${scaled_positions[i]}" -eq "$h" ]; then
                line+="*"
            else
                line+=" " # Otherwise, print a space.
            fi
        done
        echo "$line"
    done

    # Draw the X-axis line.
    printf "%6s +%s\n" " " "$(printf -- "-%.0s" $(seq 1 "$num_data_points"))"

    # Draw X-axis labels (time).
    printf "%6s  " " " # Indent for the labels.
    local first_timestamp=${timestamps[0]}
    local last_timestamp=${timestamps[num_data_points-1]}

    local start_time_str=$(date -d "@$first_timestamp" +%H:%M) # Format first timestamp
    local end_time_str=$(date -d "@$last_timestamp" +%H:%M)   # Format last timestamp

    local mid_time_str=""
    local mid_len=0

    # Only show a middle label if there's enough horizontal space.
    if [ "$num_data_points" -gt 20 ]; then
        local mid_timestamp_index=$((num_data_points / 2))
        mid_time_str=$(date -d "@${timestamps[$mid_timestamp_index]}" +%H:%M)
        mid_len=${#mid_time_str}
    fi

    local start_len=${#start_time_str}
    local end_len=${#end_time_str}
    local total_label_width=$((start_len + mid_len + end_len))

    # Distribute labels based on available space.
    if [ "$total_label_width" -le "$num_data_points" ]; then
        # Enough space for all three labels (start, middle, end).
        local padding_before_mid=$(( (num_data_points - total_label_width) / 2 ))
        local padding_after_mid=$(( num_data_points - total_label_width - padding_before_mid ))

        printf "%s%*s%s%*s%s\n" \
            "$start_time_str" \
            "$padding_before_mid" " " \
            "$mid_time_str" \
            "$padding_after_mid" " " \
            "$end_time_str"
    else
        # Not enough space, just show start and end labels.
        local padding_between=$((num_data_points - start_len - end_len))
        if [ "$padding_between" -lt 0 ]; then padding_between=0; fi # Ensure padding is not negative.
        printf "%s%*s%s\n" "$start_time_str" "$padding_between" " " "$end_time_str"
    fi

    echo "" # Empty line for spacing.
    echo "Current Chrome Memory: $(get_chrome_memory_usage) MB"
    echo "Press Ctrl+C to exit."
}

# --- Main Loop ---

# This loop runs indefinitely until interrupted (e.g., by Ctrl+C).
while true; do
    tput clear # Clear the terminal screen before redrawing the chart.

    # 1. Get current Chrome memory usage.
    current_mem=$(get_chrome_memory_usage)
    current_timestamp=$(date +%s)

    # 2. Append the current timestamp and memory usage to the data file.
    echo "$current_timestamp:$current_mem" >> "$DATA_FILE"

    # 3. Clean up old data from the file to keep only the last DURATION_HOURS
    #    and ensure it fits the terminal width.
    clean_old_data

    # 4. Draw the ASCII chart based on the current data.
    draw_chart

    # 5. Wait for the specified interval before the next update.
    sleep "$INTERVAL_SECONDS"
done