#!/bin/bash
##########################################################################
# 利用例:
# ssm-run-command.sh
# ssm-run-command.sh -t env:tes -d AWS-RunShellScript -f parameters.json
#
# オプション:
#   -t <タグ>            # タグを指定 (key:value形式)
#   -d <ドキュメント名>    # 実行するSSM Run Commandドキュメント名を指定
#   -f <パラメータファイル> # SSM Run CommandドキュメントのパラメータをJSONファイルで指定
##########################################################################

TAG=""
DOCUMENT_NAME=""
PARAMETERS_FILE=""

while getopts "t:d:f:" opt; do
  case $opt in
    t)
      TAG="$OPTARG"
      ;;
    d)
      DOCUMENT_NAME="$OPTARG"
      ;;
    f)
      PARAMETERS_FILE="$OPTARG"
      ;;
    \?)
      echo "無効なオプション: -$OPTARG" >&2
      echo "Usage: $0 [-t タグ] [-d ドキュメント名] [-f パラメータファイル]"
      exit 1
      ;;
    :)
      echo "オプション -$OPTARG には引数が必要です。" >&2
      exit 1
      ;;
  esac
done

CACHE_DIR="/tmp/ssm_cache_$$"
mkdir -p "$CACHE_DIR"

cleanup() {
  if [[ -d "$CACHE_DIR" ]]; then
    rm "$CACHE_DIR"/*.json
    rmdir "$CACHE_DIR"
  fi
}
trap cleanup EXIT

# ssmが認識しているインスタンスのID一覧取得
get_ssm_instances() {
  aws ssm describe-instance-information --query "InstanceInformationList[*].InstanceId" --output text
}

# キャッシュファイル作成
cache_instance_info() {
  local instance_ids=("$@")

  aws ec2 describe-instances \
    --instance-ids "${instance_ids[@]}" \
    --query "Reservations[].Instances[].[InstanceId,Tags]" \
    --output json > "$CACHE_DIR/instance_data.json"

  jq -c '.[]' "$CACHE_DIR/instance_data.json" | while read -r line; do
    local instance_id
    instance_id=$(echo "$line" | jq -r '.[0]')
    echo "$line" | jq -r '.[1]' > "$CACHE_DIR/${instance_id}.json"
  done
}

# インスタンスIDからNameタグ取得
get_instance_name() {
  local instance_id="$1"
  local cache_file="$CACHE_DIR/${instance_id}.json"

  if [[ -f "$cache_file" ]]; then
    jq -r 'map(select(.Key == "Name") | .Value) | .[0] // "N/A"' "$cache_file"
  else
    echo "N/A"
  fi
}

# インスタンスIDからタグ一覧取得
get_instance_tags() {
  local instance_id="$1"
  local cache_file="$CACHE_DIR/${instance_id}.json"

  if [[ -f "$cache_file" ]]; then
    jq -r 'if . == null then "" else (map("\(.Key):\(.Value)") | join(" ")) end' "$cache_file"
  else
    echo ""
  fi
}

# ユーザー入力確認プロンプト
prompt_confirmation() {
  local prompt_message="$1"
  read -rp "$prompt_message (y/n): " answer
  if [[ "$answer" != "y" && "$answer" != "Y" ]]; then
    echo "キャンセルされました。"
    exit 0
  fi
}

# インスタンス一覧表示
display_instance_list() {
  local instances=("$@")
  local index=1

  for inst in "${instances[@]}"; do
    name=$(get_instance_name "$inst")
    tags=$(get_instance_tags "$inst")

    echo "$index) Name: ${name} / InstanceID: ${inst} / Tags: ${tags}"
    ((index++))
  done
}

# 指定されたタグを含むインスタンスを検索
filter_instances_by_tag() {
  local instances=("$@")
  local filtered_instances=()

  for inst in "${instances[@]}"; do
    tags=$(get_instance_tags "$inst")
    if [[ "$tags" == *"$TAG"* ]]; then
      filtered_instances+=("$inst")
    fi
  done

  echo "${filtered_instances[@]}"
}

run_ssm_command() {
  local instance_ids=("$@")

  cmd="aws ssm send-command \
    --instance-ids ${instance_ids[*]} \
    --document-name \"$DOCUMENT_NAME\" \
    --cli-input-json file://$PARAMETERS_FILE \
    --query \"Command.CommandId\" --output text"

  command_id=$(eval "$cmd")
  if [[ $? -ne 0 ]];then
    exit
  fi

  echo ""
  echo "Command ID: ${command_id}"
  echo "コマンドを実行中..."

  local remaining_instances=("${instance_ids[@]}")

  while [[ ${#remaining_instances[@]} -gt 0 ]]; do
    local next_round_instances=()

    for instance in "${remaining_instances[@]}"; do
      status=$(aws ssm list-command-invocations \
        --command-id "$command_id" \
        --instance-id "$instance" \
        --query "CommandInvocations[0].Status" --output text 2>/dev/null)

      if [[ "$status" == "Success" || "$status" == "Failed" || "$status" == "Cancelled" || "$status" == "TimedOut" ]]; then
        instance_name=$(get_instance_name "$instance")
        echo "----------------------------------------"
        echo "Name: ${instance_name} / Instance: ${instance} / Status: ${status}"
        echo ""
        echo "Log:"
        output=$(aws ssm list-command-invocations \
          --command-id "$command_id" \
          --instance-id "$instance" \
          --details \
          --query "CommandInvocations[0].CommandPlugins[0].Output" --output text)
        echo "$output"
        echo ""
      else
        next_round_instances+=("$instance")
      fi
    done

    if [[ ${#next_round_instances[@]} -gt 0 ]]; then
      echo "----------------------------------------"
      echo "未完了のインスタンスあり。5秒待機して再確認..."
      sleep 5
    fi

    remaining_instances=("${next_round_instances[@]}")
  done

  echo "----------------------------------------"
  echo "すべてのインスタンスでコマンドが完了しました。"
}

main() {
  echo ""
  full_instances=($(get_ssm_instances))
  if [[ ${#full_instances[@]} -eq 0 ]]; then
    echo "SSM に認識されているインスタンスがありません。"
    exit 1
  fi

  cache_instance_info "${full_instances[@]}"
  if [[ -z "$TAG" ]]; then
    echo "[SSM で認識されている全インスタンス]"
    display_instance_list "${full_instances[@]}"
    echo ""
    while [[ -z "$TAG" ]]; do
      read -rp "タグを入力して下さい。（例: Name:hogehoge）: " TAG
    done;
  fi

  filtered_instances=($(filter_instances_by_tag "${full_instances[@]}"))
  if [[ ${#filtered_instances[@]} -eq 0 ]]; then
    echo "指定したタグが付与されたインスタンスがありません。"
    exit 1
  fi

  while [[ -z "$DOCUMENT_NAME" ]]; do
    read -rp "ドキュメント名を入力して下さい。（例: AWS-RunShellScript）: " DOCUMENT_NAME
  done;

  while [[ -z "$PARAMETERS_FILE" ]]; do
    read -rp "パラメータファイルのパスを入力して下さい。（例: parameters.json）: " PARAMETERS_FILE
  done;

  if [[ ! -f "$PARAMETERS_FILE" ]]; then
    echo "エラー: 指定されたパラメータファイル '$PARAMETERS_FILE' が存在しません。"
    exit 1
  fi

  echo ""
  echo "----------------------------------------"
  echo "[実行対象のインスタンス]"
  display_instance_list "${filtered_instances[@]}"
  echo ""
  echo "----------------------------------------"
  echo "[ドキュメント]: $DOCUMENT_NAME"
  echo ""
  echo "----------------------------------------"
  echo "[パラメータファイル]: $PARAMETERS_FILE"
  cat "$PARAMETERS_FILE"
  echo ""
  echo "----------------------------------------"
  echo ""
  prompt_confirmation "上記の内容で実行してもよろしいですか？"

  run_ssm_command "${filtered_instances[@]}"
}

main