#!/usr/bin/env bash
set -euo pipefail

# K8Studio AWS EKS SSO lab
#
# This script helps configure an existing EKS cluster for K8Studio access
# through AWS IAM Identity Center / AWS SSO and kubeconfig exec auth.
#
# It does not create or delete an EKS cluster.
#
# Requirements:
# - aws CLI v2
# - kubectl
# - An AWS admin/profile that can manage IAM Identity Center and EKS access entries
# - IAM Identity Center enabled
# - An existing EKS cluster
#
# Usage:
#   ./k8studio-eks-sso-lab.sh discover
#   ./k8studio-eks-sso-lab.sh permission-set
#   ./k8studio-eks-sso-lab.sh assign
#   ./k8studio-eks-sso-lab.sh write-profile
#   ./k8studio-eks-sso-lab.sh login
#   ./k8studio-eks-sso-lab.sh eks-access
#   ./k8studio-eks-sso-lab.sh kubeconfig
#   ./k8studio-eks-sso-lab.sh verify
#   ./k8studio-eks-sso-lab.sh cleanup-access

AWS_ADMIN_PROFILE="${AWS_ADMIN_PROFILE:-cloudops}"
AWS_SSO_PROFILE="${AWS_SSO_PROFILE:-k8studio-sso}"
AWS_REGION="${AWS_REGION:-us-east-1}"
AWS_ACCOUNT_ID="${AWS_ACCOUNT_ID:-}"
SSO_START_URL="${SSO_START_URL:-}"
SSO_REGION="${SSO_REGION:-us-east-1}"
SSO_SESSION_NAME="${SSO_SESSION_NAME:-k8studio-sso}"
SSO_ROLE_NAME="${SSO_ROLE_NAME:-K8StudioEKSExec}"
PERMISSION_SET_NAME="${PERMISSION_SET_NAME:-K8StudioEKSExec}"
PRINCIPAL_TYPE="${PRINCIPAL_TYPE:-USER}"
PRINCIPAL_ID="${PRINCIPAL_ID:-}"
CLUSTER_NAME="${CLUSTER_NAME:-k8studio-auth-eks}"
EKS_ACCESS_POLICY_ARN="${EKS_ACCESS_POLICY_ARN:-arn:aws:eks::aws:cluster-access-policy/AmazonEKSClusterAdminPolicy}"
KUBECONFIG_PATH="${KUBECONFIG_PATH:-${HOME}/.kube/k8studio-eks-sso.yaml}"

usage() {
  cat <<USAGE
K8Studio AWS EKS SSO lab

Commands:
  discover          Print IAM Identity Center instance details
  permission-set    Create/update the IAM Identity Center permission set
  assign            Assign a USER or GROUP to the AWS account
  write-profile     Write the local AWS SSO profile to ~/.aws/config
  login             Run aws sso login with the SSO profile
  eks-access        Create EKS access entry and associate access policy
  kubeconfig        Generate an SSO-backed kubeconfig
  verify            Verify kubectl can reach the cluster
  cleanup-access    Remove EKS access entry and account assignment

Required environment for most commands:
  AWS_ACCOUNT_ID          AWS account ID that contains the EKS cluster
  SSO_START_URL           IAM Identity Center start URL
  PRINCIPAL_ID            Identity Store user ID or group ID for assignment

Common environment:
  AWS_ADMIN_PROFILE=${AWS_ADMIN_PROFILE}
  AWS_SSO_PROFILE=${AWS_SSO_PROFILE}
  AWS_REGION=${AWS_REGION}
  SSO_REGION=${SSO_REGION}
  SSO_SESSION_NAME=${SSO_SESSION_NAME}
  SSO_ROLE_NAME=${SSO_ROLE_NAME}
  PERMISSION_SET_NAME=${PERMISSION_SET_NAME}
  PRINCIPAL_TYPE=${PRINCIPAL_TYPE}
  CLUSTER_NAME=${CLUSTER_NAME}
  EKS_ACCESS_POLICY_ARN=${EKS_ACCESS_POLICY_ARN}
  KUBECONFIG_PATH=${KUBECONFIG_PATH}

Example:
  export AWS_ACCOUNT_ID=428380785741
  export SSO_START_URL=https://example.awsapps.com/start
  export PRINCIPAL_ID=xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx

  ./k8studio-eks-sso-lab.sh discover
  ./k8studio-eks-sso-lab.sh permission-set
  ./k8studio-eks-sso-lab.sh assign
  ./k8studio-eks-sso-lab.sh write-profile
  ./k8studio-eks-sso-lab.sh login
  ./k8studio-eks-sso-lab.sh eks-access
  ./k8studio-eks-sso-lab.sh kubeconfig
  ./k8studio-eks-sso-lab.sh verify
USAGE
}

require_command() {
  if ! command -v "$1" >/dev/null 2>&1; then
    echo "Missing required command: $1" >&2
    exit 1
  fi
}

require_env() {
  local name="$1"
  if [[ -z "${!name:-}" ]]; then
    echo "Missing required environment variable: ${name}" >&2
    exit 1
  fi
}

aws_admin() {
  aws --profile "${AWS_ADMIN_PROFILE}" --region "${AWS_REGION}" "$@"
}

identity_center_instance() {
  aws_admin sso-admin list-instances \
    --query 'Instances[0].{InstanceArn:InstanceArn,IdentityStoreId:IdentityStoreId}' \
    --output json
}

instance_arn() {
  identity_center_instance | node -e "let data=''; process.stdin.on('data', c => data += c); process.stdin.on('end', () => console.log(JSON.parse(data).InstanceArn || ''));"
}

identity_store_id() {
  identity_center_instance | node -e "let data=''; process.stdin.on('data', c => data += c); process.stdin.on('end', () => console.log(JSON.parse(data).IdentityStoreId || ''));"
}

permission_set_arn() {
  local arn
  arn="$(aws_admin sso-admin list-permission-sets \
    --instance-arn "$(instance_arn)" \
    --query 'PermissionSets[]' \
    --output text | tr '\t' '\n' | while read -r candidate; do
      [[ -z "${candidate}" ]] && continue
      name="$(aws_admin sso-admin describe-permission-set \
        --instance-arn "$(instance_arn)" \
        --permission-set-arn "${candidate}" \
        --query 'PermissionSet.Name' \
        --output text)"
      if [[ "${name}" == "${PERMISSION_SET_NAME}" ]]; then
        echo "${candidate}"
        break
      fi
    done)"
  echo "${arn}"
}

sso_role_arn() {
  require_env AWS_ACCOUNT_ID

  local role_name
  role_name="$(aws_admin iam list-roles \
    --query "Roles[?starts_with(RoleName, 'AWSReservedSSO_${SSO_ROLE_NAME}_')].RoleName | [0]" \
    --output text)"

  if [[ -z "${role_name}" || "${role_name}" == "None" ]]; then
    echo "Could not find AWSReservedSSO role for ${SSO_ROLE_NAME}. Did you run assign?" >&2
    exit 1
  fi

  aws_admin iam get-role --role-name "${role_name}" --query 'Role.Arn' --output text
}

discover() {
  require_command aws
  echo "IAM Identity Center instance:"
  identity_center_instance

  echo
  echo "Identity Store ID:"
  identity_store_id
}

create_permission_set() {
  require_command aws

  local instance permission_set
  instance="$(instance_arn)"
  permission_set="$(permission_set_arn)"

  if [[ -z "${permission_set}" ]]; then
    permission_set="$(aws_admin sso-admin create-permission-set \
      --instance-arn "${instance}" \
      --name "${PERMISSION_SET_NAME}" \
      --session-duration PT8H \
      --query 'PermissionSet.PermissionSetArn' \
      --output text)"
  fi

  aws_admin sso-admin put-inline-policy-to-permission-set \
    --instance-arn "${instance}" \
    --permission-set-arn "${permission_set}" \
    --inline-policy '{
      "Version": "2012-10-17",
      "Statement": [
        {
          "Effect": "Allow",
          "Action": [
            "eks:DescribeCluster",
            "sts:GetCallerIdentity"
          ],
          "Resource": "*"
        }
      ]
    }'

  echo "Permission set ready:"
  echo "${permission_set}"
}

assign_account() {
  require_command aws
  require_env AWS_ACCOUNT_ID
  require_env PRINCIPAL_ID

  local instance permission_set
  instance="$(instance_arn)"
  permission_set="$(permission_set_arn)"

  if [[ -z "${permission_set}" ]]; then
    echo "Permission set does not exist. Run: $0 permission-set" >&2
    exit 1
  fi

  aws_admin sso-admin create-account-assignment \
    --instance-arn "${instance}" \
    --target-id "${AWS_ACCOUNT_ID}" \
    --target-type AWS_ACCOUNT \
    --permission-set-arn "${permission_set}" \
    --principal-type "${PRINCIPAL_TYPE}" \
    --principal-id "${PRINCIPAL_ID}" >/dev/null

  echo "Account assignment requested."
  echo "AWS may take a minute to create the AWSReservedSSO role."
}

write_profile() {
  require_env AWS_ACCOUNT_ID
  require_env SSO_START_URL

  mkdir -p "${HOME}/.aws"
  touch "${HOME}/.aws/config"

  local backup
  backup="${HOME}/.aws/config.k8studio-backup-$(date +%Y%m%d%H%M%S)"
  cp "${HOME}/.aws/config" "${backup}"

  cat >> "${HOME}/.aws/config" <<PROFILE

[profile ${AWS_SSO_PROFILE}]
sso_session = ${SSO_SESSION_NAME}
sso_account_id = ${AWS_ACCOUNT_ID}
sso_role_name = ${SSO_ROLE_NAME}
region = ${AWS_REGION}
output = json

[sso-session ${SSO_SESSION_NAME}]
sso_start_url = ${SSO_START_URL}
sso_region = ${SSO_REGION}
sso_registration_scopes = sso:account:access
PROFILE

  echo "Wrote AWS SSO profile: ${AWS_SSO_PROFILE}"
  echo "Backup: ${backup}"
}

login() {
  require_command aws
  aws sso login --profile "${AWS_SSO_PROFILE}"
}

create_eks_access() {
  require_command aws

  local principal
  principal="$(sso_role_arn)"

  aws_admin eks create-access-entry \
    --cluster-name "${CLUSTER_NAME}" \
    --principal-arn "${principal}" >/dev/null 2>&1 || true

  aws_admin eks associate-access-policy \
    --cluster-name "${CLUSTER_NAME}" \
    --principal-arn "${principal}" \
    --policy-arn "${EKS_ACCESS_POLICY_ARN}" \
    --access-scope type=cluster >/dev/null 2>&1 || true

  echo "EKS access entry ready:"
  echo "${principal}"
}

write_kubeconfig() {
  require_command aws
  mkdir -p "$(dirname "${KUBECONFIG_PATH}")"

  aws eks update-kubeconfig \
    --name "${CLUSTER_NAME}" \
    --region "${AWS_REGION}" \
    --profile "${AWS_SSO_PROFILE}" \
    --kubeconfig "${KUBECONFIG_PATH}"

  echo "Kubeconfig written:"
  echo "${KUBECONFIG_PATH}"
}

verify() {
  require_command kubectl
  kubectl --kubeconfig "${KUBECONFIG_PATH}" get nodes
}

cleanup_access() {
  require_command aws
  require_env AWS_ACCOUNT_ID

  local principal instance permission_set
  principal="$(sso_role_arn)"
  instance="$(instance_arn)"
  permission_set="$(permission_set_arn)"

  aws_admin eks disassociate-access-policy \
    --cluster-name "${CLUSTER_NAME}" \
    --principal-arn "${principal}" \
    --policy-arn "${EKS_ACCESS_POLICY_ARN}" >/dev/null 2>&1 || true

  aws_admin eks delete-access-entry \
    --cluster-name "${CLUSTER_NAME}" \
    --principal-arn "${principal}" >/dev/null 2>&1 || true

  if [[ -n "${permission_set}" && -n "${PRINCIPAL_ID}" ]]; then
    aws_admin sso-admin delete-account-assignment \
      --instance-arn "${instance}" \
      --target-id "${AWS_ACCOUNT_ID}" \
      --target-type AWS_ACCOUNT \
      --permission-set-arn "${permission_set}" \
      --principal-type "${PRINCIPAL_TYPE}" \
      --principal-id "${PRINCIPAL_ID}" >/dev/null 2>&1 || true
  fi

  echo "Removed EKS access entry. Removed account assignment if PRINCIPAL_ID was provided."
}

case "${1:-}" in
  discover)
    discover
    ;;
  permission-set)
    create_permission_set
    ;;
  assign)
    assign_account
    ;;
  write-profile)
    write_profile
    ;;
  login)
    login
    ;;
  eks-access)
    create_eks_access
    ;;
  kubeconfig)
    write_kubeconfig
    ;;
  verify)
    verify
    ;;
  cleanup-access)
    cleanup_access
    ;;
  -h|--help|help|"")
    usage
    ;;
  *)
    echo "Unknown command: $1" >&2
    usage
    exit 1
    ;;
esac