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}"

# 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

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

# Get count and last usage line number for a profile
get_frequency_info() {
    local profile="$1"
    local count=0
    local last_line=0
    if [[ -f "$HISTORY_FILE" ]]; then
        count=$(grep -cFx "$profile" "$HISTORY_FILE" 2>/dev/null) || true
        last_line=$(grep -nFx "$profile" "$HISTORY_FILE" 2>/dev/null | tail -1 | cut -d: -f1) || true
    fi
    echo "${count:-0} ${last_line:-0}"
}

# Sort profiles by frequency (most used first), then recency, then alphabetical
sort_by_frequency() {
    local profiles="$1"
    if [[ ! -f "$HISTORY_FILE" ]]; then
        echo "$profiles"
        return
    fi

    # Sort by count desc, then last usage desc, then name asc
    while IFS= read -r profile; do
        read -r count last_line <<< "$(get_frequency_info "$profile")"
        printf '%d\t%d\t%s\n' "$count" "$last_line" "$profile"
    done <<< "$profiles" | sort -t$'\t' -k1,1rn -k2,2rn -k3,3 | cut -f3
}

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

profiles=$(get_profiles)
profiles=$(sort_by_frequency "$profiles")

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"
    elif command -v rofi &>/dev/null; then
        rofi -dmenu -p "$prompt"
    elif command -v dmenu &>/dev/null; then
        dmenu -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

case "$action" in
    "Console (browser)")
        exec assume -c "$selected"
        ;;
    "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