assume-menu

#!/usr/bin/env bash
# assume-menu - A dmenu-like launcher for AWS assume
# Works with fuzzel (with icons), wofi, rofi, or dmenu

set -euo pipefail

# Terminal to use for shell sessions
TERMINAL="${ASSUME_TERMINAL:-alacritty}"

# AWS config file
AWS_CONFIG="${AWS_CONFIG_FILE:-$HOME/.aws/config}"

# History file for frequency sorting
HISTORY_FILE="${ASSUME_HISTORY_FILE:-$HOME/.cache/assume-menu-history}"

# Cache file for sorted profiles
CACHE_FILE="${ASSUME_CACHE_FILE:-$HOME/.cache/assume-menu-cache}"

# Service history file (stores profile:service pairs)
SERVICE_HISTORY_FILE="${ASSUME_SERVICE_HISTORY:-$HOME/.cache/assume-menu-service-history}"

# Ensure cache directory exists
mkdir -p "$(dirname "$HISTORY_FILE")"

if [[ ! -f "$AWS_CONFIG" ]]; then
    notify-send -u critical "assume-menu" "AWS config not found at $AWS_CONFIG" 2>/dev/null || true
    echo "AWS config not found at $AWS_CONFIG" >&2
    exit 1
fi

# Compute cache key from config and history file states
get_cache_key() {
    local config_hash history_hash
    config_hash=$(md5sum "$AWS_CONFIG" 2>/dev/null | cut -d' ' -f1) || config_hash="none"
    if [[ -f "$HISTORY_FILE" ]]; then
        history_hash=$(md5sum "$HISTORY_FILE" 2>/dev/null | cut -d' ' -f1) || history_hash="none"
    else
        history_hash="empty"
    fi
    echo "${config_hash}:${history_hash}"
}

# Check if cache is valid and return cached profiles
get_cached_profiles() {
    if [[ ! -f "$CACHE_FILE" ]]; then
        return 1
    fi
    local cached_key current_key
    cached_key=$(head -1 "$CACHE_FILE")
    current_key=$(get_cache_key)
    if [[ "$cached_key" == "$current_key" ]]; then
        tail -n +2 "$CACHE_FILE"
        return 0
    fi
    return 1
}

# Save profiles to cache
save_to_cache() {
    local profiles="$1"
    {
        get_cache_key
        echo "$profiles"
    } > "$CACHE_FILE"
}

# Extract profile names from AWS config
get_profiles() {
    grep -E '^\[profile |^\[default\]' "$AWS_CONFIG" \
        | sed 's/\[profile //' \
        | sed 's/\[//' \
        | sed 's/\]//' \
        | sort
}

# Sort profiles by frequency (most used first), then recency, then alphabetical
# Reads history file once with awk instead of spawning greps per profile
sort_by_frequency() {
    local profiles="$1"
    if [[ ! -f "$HISTORY_FILE" ]]; then
        echo "$profiles"
        return
    fi

    # Use awk to compute counts and last line in a single pass
    awk '
    NR == FNR {
        # First file: history - count occurrences and track last line
        count[$0]++
        last[$0] = NR
        next
    }
    {
        # Second file: profiles - output with frequency info
        printf "%d\t%d\t%s\n", count[$0]+0, last[$0]+0, $0
    }
    ' "$HISTORY_FILE" - <<< "$profiles" | sort -t$'\t' -k1,1rn -k2,2rn -k3,3 | cut -f3
}

# Record a selection to history
record_selection() {
    echo "$1" >> "$HISTORY_FILE"
}

# Record a service selection for a profile
record_service_selection() {
    local profile="$1" service="$2"
    echo "${profile}:${service}" >> "$SERVICE_HISTORY_FILE"
}

# Sort services by frequency for a specific profile
sort_services_for_profile() {
    local profile="$1" services="$2"
    if [[ ! -f "$SERVICE_HISTORY_FILE" ]]; then
        echo "$services"
        return
    fi

    # Filter history for this profile, then sort services by frequency
    awk -v profile="$profile" '
    NR == FNR {
        # First file: service history - filter by profile and count
        if (index($0, profile ":") == 1) {
            service = substr($0, length(profile) + 2)
            count[service]++
            last[service] = NR
        }
        next
    }
    {
        # Second file: services - output with frequency info
        printf "%d\t%d\t%s\n", count[$0]+0, last[$0]+0, $0
    }
    ' "$SERVICE_HISTORY_FILE" - <<< "$services" | sort -t$'\t' -k1,1rn -k2,2rn -k3,3 | cut -f3
}

# Try to use cached profiles first
if ! profiles=$(get_cached_profiles); then
    profiles=$(get_profiles)
    profiles=$(sort_by_frequency "$profiles")
    save_to_cache "$profiles"
fi

if [[ -z "$profiles" ]]; then
    notify-send -u critical "assume-menu" "No AWS profiles found" 2>/dev/null || true
    echo "No AWS profiles found in $AWS_CONFIG" >&2
    exit 1
fi

# Format entries with icons for fuzzel (rofi dmenu protocol)
format_with_icons() {
    local icon="${1:-cloud}"
    while IFS= read -r line; do
        printf '%s\0icon\x1f%s\n' "$line" "$icon"
    done
}

# Run menu with given prompt
run_menu() {
    local prompt="$1"
    if [[ -n "${ASSUME_MENU_CMD:-}" ]]; then
        eval "$ASSUME_MENU_CMD"
    elif command -v fuzzel &>/dev/null; then
        format_with_icons "${2:-cloud}" | fuzzel --dmenu --prompt "$prompt" --width 80
    elif command -v wofi &>/dev/null; then
        wofi --dmenu --prompt "$prompt" --insensitive
    elif command -v rofi &>/dev/null; then
        rofi -dmenu -i -p "$prompt"
    elif command -v dmenu &>/dev/null; then
        dmenu -i -p "$prompt"
    else
        echo "No menu program found. Install fuzzel, wofi, rofi, or dmenu." >&2
        exit 1
    fi
}

# Step 1: Select profile
selected=$(echo "$profiles" | run_menu "profile > " "cloud") || exit 0
[[ -z "$selected" ]] && exit 0

# Record selection for frequency sorting
record_selection "$selected"

# Step 2: Select action
actions="Console (browser)
Terminal (shell)"

action=$(echo "$actions" | run_menu "action > " "utilities-terminal") || exit 0
[[ -z "$action" ]] && exit 0

shopt -s nocasematch
case "$action" in
    "Console (browser)")
        # Step 3: Select AWS service (names match granted's service_map.go)
        all_services="console
acm
apigateway
appsync
athena
backup
bedrock
billing
cloudformation
cloudfront
cloudmap
cloudwatch
codeartifact
codecommit
codedeploy
codepipeline
codesuite
cognito
config
controltower
directconnect
dms
dynamodb
ebs
ec2
ecr
ecs
eks
elasticache
elasticbeanstalk
eventbridge
globalaccelerator
grafana
iam
kms
lambda
mwaa
organizations
rds
redshift
route53
s3
sagemaker
secretsmanager
securityhub
ses
sns
sqs
ssm
stepfunctions
vpc
wafv2"

        # Sort services by history, but keep "console" always first
        sorted_services=$(sort_services_for_profile "$selected" "$all_services" | grep -v '^console$')
        services="console"$'\n'"$sorted_services"
        service=$(echo "$services" | run_menu "service > " "applications-internet") || exit 0
        [[ -z "$service" ]] && exit 0

        record_service_selection "$selected" "$service"

        if [[ "$service" == "console" ]]; then
            exec assume -c "$selected"
        else
            exec assume -c "$selected" -s "$service"
        fi
        ;;
    "Terminal (shell)")
        # Launch terminal with user's shell, run assume, then spawn interactive shell with creds
        exec "$TERMINAL" -e "${SHELL:-zsh}" -ic "assume '$selected'; ${SHELL:-zsh}"
        ;;
esac