#!/bin/bash

# Show RAM usage for containers
# Usage: ./show-ram.sh [node-path]
# Without argument: shows RAM per node and total server RAM
# With argument: shows containers for specific node

BASEPATH="$(dirname "$0")"
NODE_PATH="$1"

# Function to convert memory string to MiB
to_mib() {
    local val="$1"
    local num=$(echo "$val" | sed 's/[^0-9.]//g')
    if echo "$val" | grep -qi "GiB"; then
        echo "$num * 1024" | bc
    elif echo "$val" | grep -qi "MiB"; then
        echo "$num"
    elif echo "$val" | grep -qi "KiB"; then
        echo "$num / 1024" | bc
    else
        echo "$num"
    fi
}

# Function to format MiB to human readable
format_size() {
    local mib="$1"
    if [ "${mib%.*}" -ge 1024 ] 2>/dev/null; then
        echo "scale=2; $mib / 1024" | bc | xargs printf "%.2f GiB"
    else
        printf "%.0f MiB" "$mib"
    fi
}

if [ -z "$NODE_PATH" ]; then
    # No argument: show RAM per node (compose project)
    echo "RAM usage per node:"
    echo "========================================"

    # Find all compose files that have running containers
    declare -A node_ram

    # Get all running containers with their compose project
    while IFS= read -r line; do
        container_id=$(echo "$line" | awk '{print $1}')
        [ -z "$container_id" ] && continue

        # Get compose project label
        project=$(docker inspect "$container_id" --format '{{index .Config.Labels "com.docker.compose.project"}}' 2>/dev/null)
        [ -z "$project" ] && continue

        # Get memory usage
        mem_usage=$(docker stats --no-stream --format "{{.MemUsage}}" "$container_id" 2>/dev/null | awk -F'/' '{print $1}')
        [ -z "$mem_usage" ] && continue

        mem_mib=$(to_mib "$mem_usage")

        # Add to node total
        if [ -z "${node_ram[$project]}" ]; then
            node_ram[$project]="$mem_mib"
        else
            node_ram[$project]=$(echo "${node_ram[$project]} + $mem_mib" | bc)
        fi
    done < <(docker ps --format "{{.ID}}")

    if [ ${#node_ram[@]} -eq 0 ]; then
        echo "No running containers found"
        exit 0
    fi

    # Sort by RAM usage and display
    total_mib=0
    for project in "${!node_ram[@]}"; do
        echo "${node_ram[$project]} $project"
        total_mib=$(echo "$total_mib + ${node_ram[$project]}" | bc)
    done | sort -rn | while read mib name; do
        printf "%-50s %s\n" "$name" "$(format_size $mib)"
    done

    echo "========================================"
    echo "Total container RAM: $(format_size $total_mib)"

    # Show server total RAM
    total_server_ram=$(free -h | awk '/^Mem:/ {print $2}')
    used_server_ram=$(free -h | awk '/^Mem:/ {print $3}')
    echo "Server RAM: ${used_server_ram} / ${total_server_ram}"
else
    # Specific node - path maps directly to compose file
    # e.g., "op/reth/base-mainnet-op-reth-archive-trace" -> "op/reth/base-mainnet-op-reth-archive-trace.yml"
    COMPOSE_FILE="${BASEPATH}/${NODE_PATH}.yml"
    COMPOSE_NAME=$(basename "$NODE_PATH")

    if [ ! -f "$COMPOSE_FILE" ]; then
        echo "Error: Compose file not found: $COMPOSE_FILE"
        exit 1
    fi

    # Get container IDs for this compose project
    CONTAINER_IDS=$(docker compose -f "$COMPOSE_FILE" ps -q 2>/dev/null)

    if [ -z "$CONTAINER_IDS" ]; then
        echo "No running containers found for $COMPOSE_NAME"
        exit 0
    fi

    echo "RAM usage for $COMPOSE_NAME:"
    echo "----------------------------------------"

    # Show stats for each container
    docker stats --no-stream --format "{{.Name}}\t{{.MemUsage}}\t{{.MemPerc}}" $CONTAINER_IDS

    echo "----------------------------------------"

    # Calculate total
    total_mib=0
    for cid in $CONTAINER_IDS; do
        mem_usage=$(docker stats --no-stream --format "{{.MemUsage}}" "$cid" 2>/dev/null | awk -F'/' '{print $1}')
        mem_mib=$(to_mib "$mem_usage")
        total_mib=$(echo "$total_mib + $mem_mib" | bc)
    done

    echo "Total: $(format_size $total_mib)"
fi