#!/usr/bin/env bash

#########################################################################################
#### aws-ami-finder.sh — Find the latest AWS AMI for a given OS type                 ####
#### Queries ec2 describe-images with pre-defined OS profiles.                       ####
#### Requires: bash, aws CLI                                                         ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./aws-ami-finder.sh --os amazon2023                                           ####
####   ./aws-ami-finder.sh --os ubuntu2204                                           ####
####   ./aws-ami-finder.sh --list                                                    ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Colors ────────────────────────────────────────────────────────────
if [[ -t 1 ]]; then
    RED='\033[0;31m'
    GREEN='\033[0;32m'
    YELLOW='\033[0;33m'
    BOLD='\033[1m'
    RESET='\033[0m'
else
    RED="" GREEN="" YELLOW="" BOLD="" RESET=""
fi

log()  { echo -e "${GREEN}[OK]${RESET}    $*"; }
warn() { echo -e "${YELLOW}[WARN]${RESET}  $*"; }
err()  { echo -e "${RED}[ERROR]${RESET} $*" >&2; }

# ── Defaults ──────────────────────────────────────────────────────────
OS_TYPE=""
MAX_RESULTS=10
REGION=""
SHOW_INSTANCES=false
INSTANCE_STATE="running"
RUNNING_FILTER=""
# ── Usage ─────────────────────────────────────────────────────────────
usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Find the latest AWS AMI for a given OS type.

Options:
  --os TYPE       OS type to query (required unless --list or --running)
  --region REGION AWS region to query (default: CLI configured region)
  --running [FLT] List running instances, optionally filtered: amazon2, rhel, windows
  --stopped [FLT] List stopped instances (same filters as --running)
  --all [FLT]     List all instances regardless of state
  --list          Show supported OS types
  --results N     Number of results to show (default: 10)
  -h, --help      Show this help

Supported OS types:
  amazon2         Amazon Linux 2
  amazon2023      Amazon Linux 2023
  rhel7           Red Hat Enterprise Linux 7
  rhel8           Red Hat Enterprise Linux 8
  rhel9           Red Hat Enterprise Linux 9
  ubuntu2004      Ubuntu 20.04 LTS
  ubuntu2204      Ubuntu 22.04 LTS
  ubuntu2404      Ubuntu 24.04 LTS
  windows2019     Windows Server 2019
  windows2022     Windows Server 2022

Examples:
  ./$(basename "$0") --os amazon2023
  ./$(basename "$0") --os ubuntu2204 --region us-east-1
  ./$(basename "$0") --running
  ./$(basename "$0") --running rhel
  ./$(basename "$0") --stopped amazon2
  ./$(basename "$0") --all
  ./$(basename "$0") --list
EOF
}

# ── GovCloud detection ────────────────────────────────────────────────
is_govcloud() {
    local r="$1"
    [[ "$r" == us-gov-* ]]
}

# ── OS Profiles ───────────────────────────────────────────────────────
# Sets OWNER and NAME_FILTER for the given OS type.
# Owner IDs differ between commercial and GovCloud partitions.
set_os_profile() {
    local os="$1"
    local govcloud="$2"  # true/false

    case "$os" in
        amazon2)
            OWNER=$( $govcloud && echo "045324592363" || echo "137112412989" )
            NAME_FILTER="amzn2-ami-*-x86_64-gp2"
            ;;
        amazon2023)
            OWNER=$( $govcloud && echo "045324592363" || echo "137112412989" )
            NAME_FILTER="al2023-ami-2023.*-x86_64"
            ;;
        rhel7)
            OWNER=$( $govcloud && echo "219670896067" || echo "309956199498" )
            NAME_FILTER="RHEL-7.*_HVM-*-x86_64-*"
            ;;
        rhel8)
            OWNER=$( $govcloud && echo "219670896067" || echo "309956199498" )
            NAME_FILTER="RHEL-8.*_HVM-*-x86_64-*"
            ;;
        rhel9)
            OWNER=$( $govcloud && echo "219670896067" || echo "309956199498" )
            NAME_FILTER="RHEL-9.*_HVM-*-x86_64-*"
            ;;
        ubuntu2004)
            OWNER=$( $govcloud && echo "513442679011" || echo "099720109477" )
            NAME_FILTER="ubuntu/images/hvm-ssd/ubuntu-focal-20.04-amd64-server-*"
            ;;
        ubuntu2204)
            OWNER=$( $govcloud && echo "513442679011" || echo "099720109477" )
            NAME_FILTER="ubuntu/images/hvm-ssd/ubuntu-jammy-22.04-amd64-server-*"
            ;;
        ubuntu2404)
            OWNER=$( $govcloud && echo "513442679011" || echo "099720109477" )
            NAME_FILTER="ubuntu/images/hvm-ssd-gp3/ubuntu-noble-24.04-amd64-server-*"
            ;;
        windows2019)
            OWNER=$( $govcloud && echo "801119661308" || echo "801119661308" )
            NAME_FILTER="Windows_Server-2019-English-Full-Base-*"
            ;;
        windows2022)
            OWNER=$( $govcloud && echo "801119661308" || echo "801119661308" )
            NAME_FILTER="Windows_Server-2022-English-Full-Base-*"
            ;;
        *)
            err "Unknown OS type: $1"
            err "Run with --list to see supported types"
            exit 1
            ;;
    esac
}

# ── List supported types ─────────────────────────────────────────────
list_types() {
    echo ""
    echo -e "${BOLD}Supported OS types:${RESET}"
    echo ""
    printf "  %-16s %s\n" "amazon2"      "Amazon Linux 2"
    printf "  %-16s %s\n" "amazon2023"   "Amazon Linux 2023"
    printf "  %-16s %s\n" "rhel7"        "Red Hat Enterprise Linux 7"
    printf "  %-16s %s\n" "rhel8"        "Red Hat Enterprise Linux 8"
    printf "  %-16s %s\n" "rhel9"        "Red Hat Enterprise Linux 9"
    printf "  %-16s %s\n" "ubuntu2004"   "Ubuntu 20.04 LTS"
    printf "  %-16s %s\n" "ubuntu2204"   "Ubuntu 22.04 LTS"
    printf "  %-16s %s\n" "ubuntu2404"   "Ubuntu 24.04 LTS"
    printf "  %-16s %s\n" "windows2019"  "Windows Server 2019"
    printf "  %-16s %s\n" "windows2022"  "Windows Server 2022"
    echo ""
}

# ── Parse args ────────────────────────────────────────────────────────
if [[ $# -eq 0 ]]; then
    usage
    exit 1
fi

while [[ $# -gt 0 ]]; do
    case "$1" in
        --os)       OS_TYPE="$2"; shift 2 ;;
        --region)   REGION="$2"; shift 2 ;;
        --running)  SHOW_INSTANCES=true; INSTANCE_STATE="running"
                    if [[ $# -gt 1 && "$2" != --* ]]; then
                        RUNNING_FILTER="$2"; shift
                    fi
                    shift ;;
        --stopped)  SHOW_INSTANCES=true; INSTANCE_STATE="stopped"
                    if [[ $# -gt 1 && "$2" != --* ]]; then
                        RUNNING_FILTER="$2"; shift
                    fi
                    shift ;;
        --all)      SHOW_INSTANCES=true; INSTANCE_STATE="running,stopped"
                    if [[ $# -gt 1 && "$2" != --* ]]; then
                        RUNNING_FILTER="$2"; shift
                    fi
                    shift ;;
        --list)     list_types; exit 0 ;;
        --results)  MAX_RESULTS="$2"; shift 2 ;;
        -h|--help)  usage; exit 0 ;;
        *)          err "Unknown option: $1"; usage; exit 1 ;;
    esac
done

# ── Preflight ─────────────────────────────────────────────────────────
if ! command -v aws &>/dev/null; then
    err "AWS CLI not found. Install it first."
    exit 1
fi

# ── Running instances mode (SSM-based) ────────────────────────────────
if [[ "$SHOW_INSTANCES" == "true" ]]; then
    REGION_ARGS=()
    [[ -n "$REGION" ]] && REGION_ARGS=(--region "$REGION")

    active_region="${REGION:-$(aws configure get region 2>/dev/null || echo 'not set')}"

    filter_label="all"
    grep_pattern=""

    case "$RUNNING_FILTER" in
        amazon2)    grep_pattern=$'(Amazon Linux\t2$|Linux/UNIX)'; filter_label="Amazon Linux 2 + Linux/UNIX" ;;
        amazon2023) grep_pattern=$'Amazon Linux\t2023'; filter_label="Amazon Linux 2023" ;;
        rhel)       grep_pattern="Red Hat Enterprise Linux"; filter_label="RHEL" ;;
        windows)    grep_pattern="Windows"; filter_label="Windows" ;;
        "")         ;;
        *)          err "Unknown filter: $RUNNING_FILTER (use: amazon2, amazon2023, rhel, windows)"; exit 1 ;;
    esac

    state_label="${INSTANCE_STATE//,/ + }"
    log "Querying ${filter_label} instances (${state_label}) in ${active_region}..."

    # Get instances (InstanceId, Name, Owner tag, State, PlatformDetails)
    ec2_data=$(aws ec2 describe-instances \
        ${REGION_ARGS[@]+"${REGION_ARGS[@]}"} \
        --filters "Name=instance-state-name,Values=${INSTANCE_STATE}" \
        --query 'Reservations[].Instances[].[InstanceId, Tags[?Key==`Name`].Value | [0], Tags[?Key==`Owner`].Value | [0], State.Name, PlatformDetails]' \
        --output text 2>&1) || {
        err "EC2 query failed:"
        echo "$ec2_data" >&2
        exit 1
    }

    if [[ -z "$ec2_data" ]]; then
        warn "No running instances found"
        exit 0
    fi

    # Build lookup: InstanceId → Name, Owner, State, PlatformDetails
    declare -A NAMES
    declare -A OWNERS
    declare -A STATES
    declare -A PLATFORMS
    while IFS=$'\t' read -r iid name owner state plat; do
        NAMES["$iid"]="$name"
        OWNERS["$iid"]="${owner:-}"
        STATES["$iid"]="$state"
        PLATFORMS["$iid"]="$plat"
    done <<< "$ec2_data"

    # Get SSM data for OS identification (best-effort)
    declare -A SSM_PLATFORM
    declare -A SSM_VERSION
    ssm_data=$(aws ssm describe-instance-information \
        ${REGION_ARGS[@]+"${REGION_ARGS[@]}"} \
        --query 'InstanceInformationList[].[InstanceId, PlatformName, PlatformVersion]' \
        --output text 2>/dev/null) || ssm_data=""

    if [[ -n "$ssm_data" ]]; then
        while IFS=$'\t' read -r iid platform version; do
            SSM_PLATFORM["$iid"]="$platform"
            SSM_VERSION["$iid"]="$version"
        done <<< "$ssm_data"
    fi

    # Build display lines, merging EC2 + SSM data
    display_lines=()
    for iid in "${!NAMES[@]}"; do
        name="${NAMES[$iid]}"
        owner="${OWNERS[$iid]:--}"
        [[ "$owner" == "None" || -z "$owner" ]] && owner="-"
        # Strip session suffix (e.g. "jallen-session-12345" → "jallen")
        owner="${owner%%-session-*}"
        if [[ -n "${SSM_PLATFORM[$iid]:-}" ]]; then
            platform="${SSM_PLATFORM[$iid]}"
            version="${SSM_VERSION[$iid]}"
        else
            platform="${PLATFORMS[$iid]}"
            version="(no SSM)"
        fi
        state="${STATES[$iid]}"
        display_lines+=("${name}"$'\t'"${iid}"$'\t'"${owner}"$'\t'"${state}"$'\t'"${platform}"$'\t'"${version}")
    done

    # Apply filter if set
    if [[ -n "$grep_pattern" ]]; then
        filtered=()
        for line in "${display_lines[@]}"; do
            if echo "$line" | grep -qP "$grep_pattern"; then
                filtered+=("$line")
            fi
        done
        display_lines=("${filtered[@]+"${filtered[@]}"}")
        if [[ ${#display_lines[@]} -eq 0 ]]; then
            warn "No ${filter_label} instances found"
            exit 0
        fi
    fi

    # Sort by name
    sorted=$(printf '%s\n' "${display_lines[@]}" | sort -t$'\t' -k1)

    # Calculate dynamic column widths from data
    col_name=4; col_iid=11; col_owner=5; col_state=5; col_plat=8; col_ver=7
    while IFS=$'\t' read -r name iid owner state platform version; do
        (( ${#name}    > col_name  )) && col_name=${#name}
        (( ${#iid}     > col_iid   )) && col_iid=${#iid}
        (( ${#owner}   > col_owner )) && col_owner=${#owner}
        (( ${#state}   > col_state )) && col_state=${#state}
        (( ${#platform}> col_plat  )) && col_plat=${#platform}
        (( ${#version} > col_ver   )) && col_ver=${#version}
    done <<< "$sorted"

    # Add padding
    col_name=$((col_name + 2))
    col_iid=$((col_iid + 2))
    col_owner=$((col_owner + 2))
    col_state=$((col_state + 2))
    col_plat=$((col_plat + 2))

    fmt="  %-${col_name}s %-${col_iid}s %-${col_owner}s %-${col_state}s %-${col_plat}s %s\n"

    # Build separator lines matching column widths
    sep_name=$(printf '%*s' "$col_name" '' | tr ' ' '-')
    sep_iid=$(printf '%*s' "$col_iid" '' | tr ' ' '-')
    sep_owner=$(printf '%*s' "$col_owner" '' | tr ' ' '-')
    sep_state=$(printf '%*s' "$col_state" '' | tr ' ' '-')
    sep_plat=$(printf '%*s' "$col_plat" '' | tr ' ' '-')
    sep_ver=$(printf '%*s' "$col_ver" '' | tr ' ' '-')

    # Display
    echo ""
    printf "  ${BOLD}%-${col_name}s %-${col_iid}s %-${col_owner}s %-${col_state}s %-${col_plat}s %s${RESET}\n" "Name" "Instance ID" "Owner" "State" "Platform" "Version"
    printf "$fmt" "$sep_name" "$sep_iid" "$sep_owner" "$sep_state" "$sep_plat" "$sep_ver"
    while IFS=$'\t' read -r name iid owner state platform version; do
        # Colorize state and pad manually (escape codes break printf %-Ns)
        state_pad=$(( col_state - ${#state} ))
        pad_str=$(printf '%*s' "$state_pad" '')
        case "$state" in
            running) printf -v state_str '%b' "${GREEN}${state}${RESET}${pad_str}" ;;
            stopped) printf -v state_str '%b' "${RED}${state}${RESET}${pad_str}" ;;
            *)       printf -v state_str "%-${col_state}s" "$state" ;;
        esac
        printf "  %-${col_name}s %-${col_iid}s %-${col_owner}s %s%-${col_plat}s %s\n" \
            "$name" "$iid" "$owner" "$state_str" "$platform" "$version"
    done <<< "$sorted"
    echo ""

    ssm_count=0
    for iid in "${!NAMES[@]}"; do
        [[ -n "${SSM_PLATFORM[$iid]:-}" ]] && ((ssm_count++)) || true
    done
    log "${#display_lines[@]} instance(s) shown (${ssm_count} identified via SSM)"
    exit 0
fi

if [[ -z "$OS_TYPE" ]]; then
    err "No OS type specified. Use --os TYPE or --running"
    exit 1
fi

# ── Query ─────────────────────────────────────────────────────────────
REGION_ARGS=()
[[ -n "$REGION" ]] && REGION_ARGS=(--region "$REGION")

active_region="${REGION:-$(aws configure get region 2>/dev/null || echo 'not set')}"
govcloud=false
is_govcloud "$active_region" && govcloud=true

set_os_profile "$OS_TYPE" "$govcloud"

$govcloud && log "GovCloud detected — using GovCloud owner IDs"
log "Querying AMIs for ${OS_TYPE} in ${active_region} (owner: ${OWNER})..."
echo ""

output=$(aws ec2 describe-images \
    ${REGION_ARGS[@]+"${REGION_ARGS[@]}"} \
    --owners "$OWNER" \
    --filters "Name=name,Values=${NAME_FILTER}" "Name=state,Values=available" \
    --query "reverse(sort_by(Images, &CreationDate))[:${MAX_RESULTS}].[ImageId, Name, Description, CreationDate]" \
    --output table 2>&1 | sed "s/DescribeImages/Available AMIs/") || {
    err "AWS CLI failed:"
    echo "$output" >&2
    exit 1
}

if [[ -z "$output" ]]; then
    warn "No AMIs found for ${OS_TYPE}"
    warn "Check your AWS region (current: $(aws configure get region 2>/dev/null || echo 'not set'))"
    exit 1
fi

echo "$output"