#!/usr/bin/env bash

#########################################################################################
#### ssm-fleet-runner.sh — Run commands across EC2 instances via AWS Systems Manager  ####
#### Supports tag-based targeting, output collection, patch scanning, and inventory   ####
#### Requires: bash 4+, aws-cli v2, jq                                              ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./ssm-fleet-runner.sh --run "uptime" --tag Environment=production             ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
AWS_REGION="${AWS_REGION:-}"
INSTANCE_IDS="${INSTANCE_IDS:-}"
SSM_TAG_KEY="${SSM_TAG_KEY:-}"
SSM_TAG_VALUE="${SSM_TAG_VALUE:-}"
SSM_TIMEOUT="${SSM_TIMEOUT:-600}"
SSM_MAX_CONCURRENCY="${SSM_MAX_CONCURRENCY:-50}"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
OUTPUT_DIR="${OUTPUT_DIR:-}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
RUN_MODE=""
COMMAND_TEXT=""
SCRIPT_FILE=""
START_TIME=""
ALL_INSTANCES="false"

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" DIM="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" DIM="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }

die() { err "$*"; exit 1; }

elapsed() {
    local end_time
    end_time=$(date +%s)
    echo "$(( end_time - START_TIME ))s"
}

# ── AWS CLI wrapper ───────────────────────────────────────────────────
aws_cmd() {
    local args=("$@")
    [[ -n "$AWS_REGION" ]] && args+=(--region "$AWS_REGION")
    verbose "aws ${args[*]}"
    aws "${args[@]}"
}

# ── Dependency check ──────────────────────────────────────────────────
check_deps() {
    for cmd in aws jq; do
        if ! command -v "$cmd" &>/dev/null; then
            die "${cmd} is required but not installed"
        fi
    done

    if ! aws sts get-caller-identity &>/dev/null; then
        die "AWS credentials not configured or expired"
    fi

    if [[ -z "$AWS_REGION" ]]; then
        AWS_REGION=$(aws configure get region 2>/dev/null || echo "")
        if [[ -z "$AWS_REGION" ]]; then
            die "AWS_REGION is required (set via env var or aws configure)"
        fi
    fi

    verbose "Using region: ${AWS_REGION}"
}

# ── Resolve target instances ─────────────────────────────────────────
resolve_targets() {
    local targets=()

    if [[ -n "$INSTANCE_IDS" ]]; then
        IFS=',' read -ra targets <<< "$INSTANCE_IDS"
        echo "${targets[@]}"
        return
    fi

    if [[ -n "$SSM_TAG_KEY" ]]; then
        aws_cmd ssm describe-instance-information \
            --filters "Key=tag:${SSM_TAG_KEY},Values=${SSM_TAG_VALUE:-*}" \
            --query 'InstanceInformationList[*].InstanceId' \
            --output text 2>/dev/null | tr '\t' ' '
        return
    fi

    if [[ "$ALL_INSTANCES" == "true" ]]; then
        aws_cmd ssm describe-instance-information \
            --query 'InstanceInformationList[*].InstanceId' \
            --output text 2>/dev/null | tr '\t' ' '
        return
    fi

    die "No targets specified. Use --instance-ids, --tag, or --all-instances"
}

# ── Get instance name ────────────────────────────────────────────────
get_instance_name() {
    local instance_id="$1"
    # shellcheck disable=SC2016
    aws_cmd ec2 describe-instances \
        --instance-ids "$instance_id" \
        --query 'Reservations[0].Instances[0].Tags[?Key==`Name`].Value | [0]' \
        --output text 2>/dev/null || echo "N/A"
}

# ══════════════════════════════════════════════════════════════════════
#                          RUN COMMAND MODE
# ══════════════════════════════════════════════════════════════════════

do_run() {
    local command_to_run="$COMMAND_TEXT"

    if [[ "$RUN_MODE" == "script" ]]; then
        if [[ ! -f "$SCRIPT_FILE" ]]; then
            die "Script file not found: ${SCRIPT_FILE}"
        fi
        command_to_run=$(cat "$SCRIPT_FILE")
    fi

    log "Resolving target instances..."
    local targets_str
    targets_str=$(resolve_targets)

    if [[ -z "$targets_str" ]]; then
        die "No SSM-managed instances found matching criteria"
    fi

    local -a target_array
    read -ra target_array <<< "$targets_str"
    local count=${#target_array[@]}
    log "Found ${count} target instance(s)"

    # Build target specification
    local target_spec
    if [[ -n "$INSTANCE_IDS" ]]; then
        target_spec="--instance-ids ${target_array[*]}"
    elif [[ -n "$SSM_TAG_KEY" ]]; then
        target_spec="--targets Key=tag:${SSM_TAG_KEY},Values=${SSM_TAG_VALUE:-*}"
    else
        target_spec="--instance-ids ${target_array[*]}"
    fi

    log "Sending command to ${count} instance(s)..."
    verbose "Command: ${command_to_run:0:100}..."

    local cmd_id
    # shellcheck disable=SC2086
    cmd_id=$(aws_cmd ssm send-command \
        --document-name "AWS-RunShellScript" \
        --parameters "commands=[\"${command_to_run}\"]" \
        $target_spec \
        --timeout-seconds "$SSM_TIMEOUT" \
        --max-concurrency "${SSM_MAX_CONCURRENCY}" \
        --query 'Command.CommandId' \
        --output text 2>/dev/null) || die "Failed to send SSM command"

    log "Command ID: ${cmd_id}"
    log "Waiting for completion (timeout: ${SSM_TIMEOUT}s)..."

    # Poll for completion
    local attempts=0
    local max_attempts=$((SSM_TIMEOUT / 5))
    local all_done="false"

    while [[ "$all_done" != "true" && "$attempts" -lt "$max_attempts" ]]; do
        sleep 5
        ((attempts++)) || true

        local status_json
        status_json=$(aws_cmd ssm list-command-invocations \
            --command-id "$cmd_id" \
            --query 'CommandInvocations[*].{Id:InstanceId,Status:Status}' \
            --output json 2>/dev/null) || continue

        local pending
        pending=$(echo "$status_json" | jq '[.[] | select(.Status == "InProgress" or .Status == "Pending")] | length')
        if [[ "$pending" -eq 0 ]]; then
            all_done="true"
        fi

        verbose "Poll ${attempts}: ${pending} still in progress"
    done

    # Collect results
    collect_results "$cmd_id" "${target_array[@]}"
}

collect_results() {
    local cmd_id="$1"
    shift
    local instances=("$@")

    local success=0 failed=0 timed_out=0

    echo ""
    for instance_id in "${instances[@]}"; do
        local result_json
        result_json=$(aws_cmd ssm get-command-invocation \
            --command-id "$cmd_id" \
            --instance-id "$instance_id" \
            --output json 2>/dev/null) || continue

        local status stdout stderr
        status=$(echo "$result_json" | jq -r '.Status')
        stdout=$(echo "$result_json" | jq -r '.StandardOutputContent // ""')
        stderr=$(echo "$result_json" | jq -r '.StandardErrorContent // ""')

        local name
        name=$(get_instance_name "$instance_id")
        [[ "$name" == "None" ]] && name="N/A"

        case "$status" in
            Success)
                echo -e "  ${GREEN}✓${RESET} ${instance_id} (${name})"
                ((success++)) || true
                ;;
            Failed)
                echo -e "  ${RED}✗${RESET} ${instance_id} (${name}) — failed"
                ((failed++)) || true
                ;;
            TimedOut)
                echo -e "  ${YELLOW}⏱${RESET} ${instance_id} (${name}) — timed out"
                ((timed_out++)) || true
                ;;
            *)
                echo -e "  ${DIM}?${RESET} ${instance_id} (${name}) — ${status}"
                ;;
        esac

        if [[ -n "$stdout" && "$OUTPUT_FORMAT" == "text" ]]; then
            # shellcheck disable=SC2001
            echo "$stdout" | sed 's/^/    /'
            echo ""
        fi

        if [[ -n "$stderr" && "$VERBOSE" == "true" ]]; then
            echo -e "    ${RED}stderr:${RESET}"
            # shellcheck disable=SC2001
            echo "$stderr" | sed 's/^/    /'
            echo ""
        fi

        # Save to output directory
        if [[ -n "$OUTPUT_DIR" ]]; then
            mkdir -p "$OUTPUT_DIR"
            echo "$stdout" > "${OUTPUT_DIR}/${instance_id}.stdout.txt"
            [[ -n "$stderr" ]] && echo "$stderr" > "${OUTPUT_DIR}/${instance_id}.stderr.txt"
        fi
    done

    if [[ "$OUTPUT_FORMAT" == "json" ]]; then
        aws_cmd ssm list-command-invocations \
            --command-id "$cmd_id" \
            --details \
            --output json 2>/dev/null
    fi

    echo ""
    log "Summary: success=${success}, failed=${failed}, timed_out=${timed_out}"
    log "Completed in $(elapsed)"
}

# ══════════════════════════════════════════════════════════════════════
#                          PATCH SCAN MODE
# ══════════════════════════════════════════════════════════════════════

do_patch_scan() {
    log "Running patch compliance scan..."
    local targets_str
    targets_str=$(resolve_targets)

    if [[ -z "$targets_str" ]]; then
        die "No SSM-managed instances found matching criteria"
    fi

    local -a target_array
    read -ra target_array <<< "$targets_str"
    log "Scanning ${#target_array[@]} instance(s) for patch compliance..."

    local target_spec
    if [[ -n "$SSM_TAG_KEY" ]]; then
        target_spec="--targets Key=tag:${SSM_TAG_KEY},Values=${SSM_TAG_VALUE:-*}"
    else
        target_spec="--instance-ids ${target_array[*]}"
    fi

    local cmd_id
    # shellcheck disable=SC2086
    cmd_id=$(aws_cmd ssm send-command \
        --document-name "AWS-RunPatchBaseline" \
        --parameters '{"Operation":["Scan"]}' \
        $target_spec \
        --timeout-seconds "$SSM_TIMEOUT" \
        --query 'Command.CommandId' \
        --output text 2>/dev/null) || die "Failed to send patch scan command"

    log "Command ID: ${cmd_id}"
    log "Waiting for scan completion..."

    sleep 10

    local attempts=0
    local max_attempts=$((SSM_TIMEOUT / 10))
    while [[ "$attempts" -lt "$max_attempts" ]]; do
        sleep 10
        ((attempts++)) || true

        local pending
        # shellcheck disable=SC2016
        pending=$(aws_cmd ssm list-command-invocations \
            --command-id "$cmd_id" \
            --query 'CommandInvocations[?Status==`InProgress` || Status==`Pending`] | length(@)' \
            --output text 2>/dev/null) || continue

        if [[ "$pending" -eq 0 ]]; then
            break
        fi
        verbose "Poll ${attempts}: ${pending} still scanning..."
    done

    # Display results
    echo ""
    printf "  ${BOLD}%-22s %-14s %s${RESET}\n" "INSTANCE" "STATUS" "DETAILS"
    printf "  %s\n" "$(printf '%.0s─' {1..55})"

    aws_cmd ssm list-command-invocations \
        --command-id "$cmd_id" \
        --query 'CommandInvocations[*].{Id:InstanceId,Status:Status,Detail:StatusDetails}' \
        --output json 2>/dev/null | jq -c '.[]' | while IFS= read -r inv; do
        local iid status detail
        iid=$(echo "$inv" | jq -r '.Id')
        status=$(echo "$inv" | jq -r '.Status')
        detail=$(echo "$inv" | jq -r '.Detail')

        local icon="$GREEN✓$RESET"
        [[ "$status" != "Success" ]] && icon="$RED✗$RESET"

        printf "  %-22s ${icon} %-12s %s\n" "$iid" "$status" "$detail"
    done

    echo ""
    log "Patch scan completed in $(elapsed)"
}

# ══════════════════════════════════════════════════════════════════════
#                          PATCH INSTALL MODE
# ══════════════════════════════════════════════════════════════════════

do_patch_install() {
    log "Running patch installation..."
    warn "This will install patches and may require reboots"

    local targets_str
    targets_str=$(resolve_targets)

    local -a target_array
    read -ra target_array <<< "$targets_str"
    log "Installing patches on ${#target_array[@]} instance(s)..."

    local target_spec
    if [[ -n "$SSM_TAG_KEY" ]]; then
        target_spec="--targets Key=tag:${SSM_TAG_KEY},Values=${SSM_TAG_VALUE:-*}"
    else
        target_spec="--instance-ids ${target_array[*]}"
    fi

    local cmd_id
    # shellcheck disable=SC2086
    cmd_id=$(aws_cmd ssm send-command \
        --document-name "AWS-RunPatchBaseline" \
        --parameters '{"Operation":["Install"],"RebootOption":["RebootIfNeeded"]}' \
        $target_spec \
        --timeout-seconds "$SSM_TIMEOUT" \
        --query 'Command.CommandId' \
        --output text 2>/dev/null) || die "Failed to send patch install command"

    log "Command ID: ${cmd_id}"
    log "Patches being installed — monitor with: aws ssm list-command-invocations --command-id ${cmd_id}"
    log "Completed in $(elapsed)"
}

# ══════════════════════════════════════════════════════════════════════
#                          INVENTORY MODE
# ══════════════════════════════════════════════════════════════════════

do_inventory() {
    log "Collecting software inventory..."

    local targets_str
    targets_str=$(resolve_targets)

    local -a target_array
    read -ra target_array <<< "$targets_str"

    echo ""
    printf "  ${BOLD}%-22s %-16s %-12s %-14s %s${RESET}\n" "INSTANCE" "PLATFORM" "AGENT_VER" "IP" "NAME"
    printf "  %s\n" "$(printf '%.0s─' {1..80})"

    for iid in "${target_array[@]}"; do
        local info_json
        info_json=$(aws_cmd ssm describe-instance-information \
            --filters "Key=InstanceIds,Values=${iid}" \
            --query 'InstanceInformationList[0]' \
            --output json 2>/dev/null) || continue

        local platform agent_ver ip_addr name
        platform=$(echo "$info_json" | jq -r '.PlatformName // "Unknown"')
        agent_ver=$(echo "$info_json" | jq -r '.AgentVersion // "?"')
        ip_addr=$(echo "$info_json" | jq -r '.IPAddress // "?"')
        name=$(get_instance_name "$iid")
        [[ "$name" == "None" ]] && name="N/A"

        printf "  %-22s %-16s %-12s %-14s %s\n" "$iid" "${platform:0:16}" "${agent_ver:0:12}" "$ip_addr" "${name:0:20}"
    done

    echo ""
    log "Inventory complete — ${#target_array[@]} instance(s)"
    log "Completed in $(elapsed)"
}

# ══════════════════════════════════════════════════════════════════════
#                          HELP
# ══════════════════════════════════════════════════════════════════════

show_help() {
    cat <<EOF
Usage: $SCRIPT_NAME [MODE] [OPTIONS]

Run commands across EC2 instances via AWS Systems Manager.

MODES:
    --run COMMAND           Run a shell command on target instances
    --script FILE           Run a script file on target instances
    --patch-scan            Scan instances for patch compliance
    --patch-install         Install patches (may reboot instances)
    --inventory             Collect instance inventory

TARGETING:
    --instance-ids IDS      Comma-separated instance IDs
    --tag KEY=VALUE         Target instances by tag
    --all-instances         Target all SSM-managed instances

OPTIONS:
    --timeout SECS          Command timeout (default: $SSM_TIMEOUT)
    --max-concurrency N     Max concurrent executions (default: $SSM_MAX_CONCURRENCY)
    --format FORMAT         Output: text (default), json, csv
    --output-dir DIR        Save per-instance output to directory
    --verbose               Debug output
    --no-color              Disable colored output
    --help, -h              Show this help

ENVIRONMENT VARIABLES:
    AWS_PROFILE             AWS CLI profile
    AWS_REGION              AWS region
    INSTANCE_IDS            Comma-separated instance IDs
    SSM_TAG_KEY             Target tag key
    SSM_TAG_VALUE           Target tag value
    SSM_TIMEOUT             Command timeout in seconds (default: 600)
    SSM_MAX_CONCURRENCY     Max concurrent executions (default: 50)
    OUTPUT_FORMAT           Output format (default: text)
    VERBOSE                 Debug output (default: false)
    COLOR                   Color mode: auto, always, never

EXAMPLES:
    # Run uptime on tagged instances
    ./$SCRIPT_NAME --run "uptime" --tag Environment=production

    # Run a script file
    ./$SCRIPT_NAME --script /opt/scripts/check.sh --instance-ids i-0abc123,i-0def456

    # Patch compliance scan
    ./$SCRIPT_NAME --patch-scan --tag PatchGroup=web

    # Collect inventory
    ./$SCRIPT_NAME --inventory --all-instances

    # Save output per instance
    ./$SCRIPT_NAME --run "df -h" --tag Role=web --output-dir /tmp/df-results
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                          MAIN
# ══════════════════════════════════════════════════════════════════════

main() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --run)              RUN_MODE="run"; COMMAND_TEXT="$2"; shift 2 ;;
            --script)           RUN_MODE="script"; SCRIPT_FILE="$2"; shift 2 ;;
            --patch-scan)       RUN_MODE="patch-scan"; shift ;;
            --patch-install)    RUN_MODE="patch-install"; shift ;;
            --inventory)        RUN_MODE="inventory"; shift ;;
            --instance-ids)     INSTANCE_IDS="$2"; shift 2 ;;
            --tag)
                local tag_pair="$2"
                SSM_TAG_KEY="${tag_pair%%=*}"
                SSM_TAG_VALUE="${tag_pair#*=}"
                shift 2
                ;;
            --all-instances)    ALL_INSTANCES="true"; shift ;;
            --timeout)          SSM_TIMEOUT="$2"; shift 2 ;;
            --max-concurrency)  SSM_MAX_CONCURRENCY="$2"; shift 2 ;;
            --format)           OUTPUT_FORMAT="$2"; shift 2 ;;
            --output-dir)       OUTPUT_DIR="$2"; shift 2 ;;
            --verbose)          VERBOSE="true"; shift ;;
            --no-color)         COLOR="never"; shift ;;
            --help|-h)          show_help; exit 0 ;;
            *)                  die "Unknown option: $1 (see --help)" ;;
        esac
    done

    setup_colors

    if [[ -z "$RUN_MODE" ]]; then err "No mode specified"; echo ""; show_help; exit 1; fi

    START_TIME=$(date +%s)

    echo ""
    echo -e "${BOLD}SSM Fleet Runner${RESET}"
    echo "Region: ${AWS_REGION:-$(aws configure get region 2>/dev/null || echo 'default')}"
    echo "Mode:   ${RUN_MODE}"
    echo "Time:   $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo ""

    check_deps

    case "$RUN_MODE" in
        run|script)    do_run ;;
        patch-scan)    do_patch_scan ;;
        patch-install) do_patch_install ;;
        inventory)     do_inventory ;;
    esac
}

main "$@"