#!/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

set -euo pipefail

BASEPATH="$(dirname "$0")"
source "$BASEPATH/.env"

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 [ -z "$num" ] || [ "$num" = "0" ]; then
        echo "0"
        return
    fi
    if echo "$val" | grep -qi "GiB"; then
        echo "$num * 1024" | bc -l | awk '{printf "%.2f", $1}'
    elif echo "$val" | grep -qi "MiB"; then
        echo "$num" | awk '{printf "%.2f", $1}'
    elif echo "$val" | grep -qi "KiB"; then
        echo "$num / 1024" | bc -l | awk '{printf "%.2f", $1}'
    else
        echo "$num" | awk '{printf "%.2f", $1}'
    fi
}

# Function to format MiB to human readable
format_size() {
    local mib="$1"
    # Handle empty or zero values
    if [ -z "$mib" ] || [ "$mib" = "0" ] || [ "$mib" = "0.00" ]; then
        echo "0 MiB"
        return
    fi
    # Compare as float using awk
    if awk "BEGIN {exit !($mib >= 1024)}"; then
        echo "scale=2; $mib / 1024" | bc -l | xargs printf "%.2f GiB"
    else
        printf "%.0f MiB" "$mib"
    fi
}

# Cache for docker stats (container_id -> mem_usage)
declare -A DOCKER_STATS_CACHE

# Cache for service -> container_id mapping
declare -A SERVICE_TO_CID

# Initialize caches - fetch all container stats once
init_docker_caches() {
    # Get all running container IDs with their service labels
    while IFS=$'\t' read -r cid service; do
        [ -z "$cid" ] && continue
        if [ -n "$service" ]; then
            SERVICE_TO_CID["$service"]="$cid"
        fi
    done < <(docker ps --format "{{.ID}}\t{{.Label \"com.docker.compose.service\"}}" 2>/dev/null || true)

    # Get stats for all containers at once (much faster than individual calls)
    # docker stats can take multiple container IDs, but we'll get all stats in one call
    if [ ${#SERVICE_TO_CID[@]} -gt 0 ]; then
        local all_cids=($(printf '%s\n' "${SERVICE_TO_CID[@]}" | sort -u))
        if [ ${#all_cids[@]} -gt 0 ]; then
            # Call docker stats once with all container IDs
            while IFS=$'\t' read -r cid mem_usage; do
                [ -z "$cid" ] && continue
                DOCKER_STATS_CACHE["$cid"]="$mem_usage"
            done < <(docker stats --no-stream --format "{{.ID}}\t{{.MemUsage}}" "${all_cids[@]}" 2>/dev/null || true)
        fi
    fi
}

# Function to get RAM for a compose file by matching service names
get_compose_ram() {
    local compose_file="$1"
    local total=0

    # Get service names defined in this compose file
    local services
    services=$(cat "$compose_file" 2>/dev/null | yaml2json - 2>/dev/null | jq -r '.services | keys[]' 2>/dev/null || echo "")
    [ -z "$services" ] && echo "0" && return

    for service in $services; do
        # Use cached container ID
        local cid="${SERVICE_TO_CID[$service]:-}"
        [ -z "$cid" ] && continue

        # Use cached memory usage
        local mem_usage="${DOCKER_STATS_CACHE[$cid]:-}"
        [ -z "$mem_usage" ] && continue

        # Extract just the used memory (before the /)
        mem_usage=$(echo "$mem_usage" | awk -F'/' '{print $1}' | xargs)
        [ -z "$mem_usage" ] && continue

        local mem_mib
        mem_mib=$(to_mib "$mem_usage")
        total=$(echo "$total + $mem_mib" | bc -l | awk '{printf "%.2f", $1}')
    done

    echo "$total"
}

# Blacklist - skip infrastructure compose files
blacklist=("drpc.yml" "drpc-free.yml" "drpc-home.yml" "base.yml" "rpc.yml" "monitoring.yml" "ftp.yml" "backup-http.yml")

is_blacklisted() {
    local part="$1"
    for word in "${blacklist[@]}"; do
        if echo "$part" | grep -qE "$word"; then
            return 0
        fi
    done
    return 1
}

if [ -z "$NODE_PATH" ]; then
    # No argument: show RAM per node using COMPOSE_FILE from .env
    echo "RAM usage per node:"
    echo "========================================"

    # Initialize caches once - this batches all docker calls
    init_docker_caches

    IFS=':' read -ra parts <<< "$COMPOSE_FILE"

    declare -A node_ram
    total_mib=0

    # Export caches to temp files for parallel processing
    temp_dir=$(mktemp -d)
    trap "rm -rf $temp_dir" EXIT
    
    # Write service->cid mapping to file
    for service in "${!SERVICE_TO_CID[@]}"; do
        echo "$service|${SERVICE_TO_CID[$service]}" >> "$temp_dir/service_cid.map"
    done
    
    # Write cid->mem_usage mapping to file
    for cid in "${!DOCKER_STATS_CACHE[@]}"; do
        echo "$cid|${DOCKER_STATS_CACHE[$cid]}" >> "$temp_dir/cid_mem.map"
    done

    # Process nodes in parallel using background jobs
    pids=()

    for part in "${parts[@]}"; do
        # Skip blacklisted files
        is_blacklisted "$part" && continue

        compose_file="${BASEPATH}/${part}"
        [ ! -f "$compose_file" ] && continue

        # Get node path (remove .yml extension)
        node_path="${part%.yml}"

        # Process in background for parallel execution
        (
            # Rebuild caches in this subprocess from temp files
            declare -A local_service_cid
            declare -A local_cid_mem
            
            while IFS='|' read -r service cid; do
                [ -n "$service" ] && [ -n "$cid" ] && local_service_cid["$service"]="$cid"
            done < "$temp_dir/service_cid.map" 2>/dev/null || true
            
            while IFS='|' read -r cid mem; do
                [ -n "$cid" ] && [ -n "$mem" ] && local_cid_mem["$cid"]="$mem"
            done < "$temp_dir/cid_mem.map" 2>/dev/null || true
            
            # Get service names from compose file
            services=$(cat "$compose_file" 2>/dev/null | yaml2json - 2>/dev/null | jq -r '.services | keys[]' 2>/dev/null || echo "")
            [ -z "$services" ] && exit 0
            
            total=0
            for service in $services; do
                cid="${local_service_cid[$service]:-}"
                [ -z "$cid" ] && continue
                
                mem_usage="${local_cid_mem[$cid]:-}"
                [ -z "$mem_usage" ] && continue
                
                mem_usage=$(echo "$mem_usage" | awk -F'/' '{print $1}' | xargs)
                [ -z "$mem_usage" ] && continue
                
                mem_mib=$(to_mib "$mem_usage")
                total=$(echo "$total + $mem_mib" | bc -l | awk '{printf "%.2f", $1}')
            done
            
            if [ "$total" != "0" ]; then
                # Sanitize node_path for use as filename (replace / with _)
                safe_node_path=$(echo "$node_path" | tr '/' '_')
                echo "$total|$node_path" > "$temp_dir/result_$safe_node_path"
            fi
        ) &
        pids+=($!)
    done

    # Wait for all background jobs to complete
    for pid in "${pids[@]}"; do
        wait "$pid" 2>/dev/null || true
    done

    # Collect results
    result_count=0
    for result_file in "$temp_dir"/result_*; do
        [ -f "$result_file" ] || continue
        IFS='|' read -r ram_mib node_path < "$result_file"
        [ -z "$ram_mib" ] || [ "$ram_mib" = "0" ] && continue
        node_ram["$node_path"]="$ram_mib"
        total_mib=$(echo "$total_mib + $ram_mib" | bc -l | awk '{printf "%.2f", $1}')
        result_count=$((result_count + 1))
    done

    if [ "$result_count" -eq 0 ]; then
        echo "No running nodes found"
        exit 0
    fi

    # Sort by RAM usage and display
    for node_path in "${!node_ram[@]}"; do
        echo "${node_ram[$node_path]} $node_path"
    done | sort -rn | while IFS=' ' read -r mib name; do
        printf "%-55s %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
    COMPOSE_FILE_PATH="${BASEPATH}/${NODE_PATH}.yml"

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

    # Initialize caches once
    init_docker_caches

    # Get service names from compose file
    services=$(cat "$COMPOSE_FILE_PATH" 2>/dev/null | yaml2json - 2>/dev/null | jq -r '.services | keys[]' 2>/dev/null || echo "")

    if [ -z "$services" ]; then
        echo "No services found in $NODE_PATH"
        exit 0
    fi

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

    total_mib=0
    for service in $services; do
        # Use cached container ID
        cid="${SERVICE_TO_CID[$service]:-}"
        [ -z "$cid" ] && continue

        # Use cached memory usage
        mem_info="${DOCKER_STATS_CACHE[$cid]:-}"
        if [ -z "$mem_info" ]; then
            # Fallback: get stats for this specific container if not in cache
            mem_info=$(docker stats --no-stream --format "{{.MemUsage}}\t{{.MemPerc}}" "$cid" 2>/dev/null || echo "")
            [ -z "$mem_info" ] && continue
        fi

        mem_usage=$(echo "$mem_info" | awk -F'\t' '{print $1}' | awk -F'/' '{print $1}' | xargs)
        mem_perc=$(echo "$mem_info" | awk -F'\t' '{print $2}' | xargs)

        # If we don't have percentage, try to get it separately
        if [ -z "$mem_perc" ]; then
            mem_perc=$(docker stats --no-stream --format "{{.MemPerc}}" "$cid" 2>/dev/null || echo "")
        fi

        printf "%-40s %s\t%s\n" "$service" "$mem_usage" "$mem_perc"

        mem_mib=$(to_mib "$mem_usage")
        total_mib=$(echo "$total_mib + $mem_mib" | bc -l | awk '{printf "%.2f", $1}')
    done

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