#!/usr/bin/env bash

set -e

# Color support

function disable_color() {
    IS_TTY=false
    txtrst=
    txtbld=
    bldred=
    bldgrn=
    bldylw=
    bldblu=
    bldmag=
    bldcyn=
}

IS_TTY=false
if [ -t 1 ]; then
    if command -v tput >/dev/null; then
        IS_TTY=true
    fi
fi
if [ "$IS_TTY" = "true" ]; then
    txtrst=$(tput sgr0 || echo '\e[0m')              # Reset
    txtbld=$(tput bold || echo '\e[1m')              # Bold
    bldred=${txtbld}$(tput setaf 1 || echo '\e[31m') # Red
    bldgrn=${txtbld}$(tput setaf 2 || echo '\e[32m') # Green
    bldylw=${txtbld}$(tput setaf 3 || echo '\e[33m') # Yellow
    bldblu=${txtbld}$(tput setaf 4 || echo '\e[34m') # Blue
    bldmag=${txtbld}$(tput setaf 5 || echo '\e[35m') # Magenta
    bldcyn=${txtbld}$(tput setaf 8 || echo '\e[38m') # Cyan
else
    disable_color
fi

# Logging

# Print the given message in cyan, but only when --verbose was passed
function debug() {
    if [ ! -z "$VERBOSE" ]; then
        printf '%s%s%s\n' "$bldcyn" "$1" "$txtrst"
    fi
}

# Print the given message in blue
function info() {
    printf '%s%s%s\n' "$bldblu" "$1" "$txtrst"
}

# Print the given message in magenta
function action() {
    printf '%s%s%s\n' "$bldmag" "$1" "$txtrst"
}

# Print the given message in yellow
function warn() {
    printf '%s%s%s\n' "$bldylw" "$1" "$txtrst"
}

# Like warn, but expects the message via redirect
function warnb() {
    printf '%s' "$bldylw"
    while read -r data; do
        printf '%s\n' "$data"
    done
    printf '%s\n' "$txtrst"
}

# Print the given message in red
function error() {
    printf '%s%s%s\n' "$bldred" "$1" "$txtrst"
    exit 1
}

# Like error, but expects the message via redirect
function errorb() {
    printf '%s' "$bldred"
    while read -r data; do
        printf '%s\n' "$data"
    done
    printf '%s\n' "$txtrst"
    exit 1
}

# Print the given message in green
function success() {
    printf '%s%s%s\n' "$bldgrn" "$1" "$txtrst"
}


# Print help if requested
function help() {
    cat << EOF
POA Infrastructure Management Tool

Usage: 

       ./infra [global options] <task> [task args]

This script will bootstrap required AWS resources, then generate infrastructure via Terraform.

Tasks:

    help              Show help
    provision         Run the provisioner to generate or modify POA infrastructure
    destroy           Tear down any provisioned resources and local state
    resources         List ARNs of any generated resources (* see docs for caveats)

Global Options:

    -v | --verbose    This will print out verbose execution information for debugging
    -h | --help       Print this help message
    --dry-run         Perform as many actions as possible without performing side-effects
    --no-color        Turn off color
    --skip-approval   Automatically accept any prompts for confirmation
    --profile=<name>  Use a specific AWS profile rather than the default
EOF
    exit 2
}

# Verify tools
function check_prereqs() {
    if ! which jq >/dev/null; then
        warnb << EOF
This script requires that the 'jq' utility has been installed and can be found in $PATH

On macOS, with Homebrew, this is as simple as 'brew install jq'.

For installs on other platforms, see https://stedolan.github.io/jq/download/
EOF
        exit 2
    fi

    if ! which aws >/dev/null; then
        warnb << EOF
This script requires that the AWS CLI tool has been installed and can be found in $PATH

On macOS, with Homebrew, this is as simple as 'brew install awscli'.

For installs on other platforms, see https://docs.aws.amazon.com/cli/latest/userguide/installing.html
EOF
        exit 2
    fi

    if ! which terraform >/dev/null; then
        warnb << EOF
  This script requires that the Terraform CLI be installed and available in PATH!

On macOS, with Homebrew, this is as simple as 'brew install terraform'.

For other platforms, see https://www.terraform.io/intro/getting-started/install.html
EOF
        exit 2
    fi
}

# Load a value which is present in one of the Terraform config
# files in the current directory, with precedence such that user-provided
# .tfvars are loaded after main.tfvars, allowing one to override those values
function get_config() {
    EXTRA_VARS="$(find . -name '*.tfvars' -and \! \( -name 'backend.tfvars' \))"
    if [ ! -z "$EXTRA_VARS" ]; then
        # shellcheck disable=SC2086 disable=2002
        cat $EXTRA_VARS | \
            grep -E "^$1 " | \
            tail -n 1 | \
            sed -r -e 's/^[^=]*= //' -e 's/"//g'
    fi
}

function destroy_bucket() {
    bucket="$(grep 'bucket' backend.tfvars | sed -e 's/bucket = //' -e 's/"//g')"
    read -r -p "Are you super sure you want to delete the Terraform state bucket and all versions? (y/n) "
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        exit 2
    fi
    # Delete all versions and delete markers first
    info "Disabling bucket versioning for S3 bucket '$bucket'.."
    aws s3api put-bucket-versioning --bucket="$bucket" --versioning-configuration="Status=Suspended"
    info "Deleting old versions of S3 bucket '$bucket'.."
    # shellcheck disable=SC1004
    aws s3api list-object-versions --bucket="$bucket" |\
        jq '.Versions[], .DeleteMarkers[] | "\"\(.Key)\" \"\(.VersionId)\""' --raw-output |\
        awk -v bucket="$bucket" '{ \
          print "aws s3api delete-object", \
              "--bucket=\"" bucket "\"", \
              "--key=\"" $1 "\"", \
              "--version-id=\"" $2 "\"" \
          | "/bin/sh >/dev/null"; \
          print "Deleted version " $2 "of " $1 " successfully"; \
         }'
    # Finally, delete the bucket and all its contents
    aws s3 rb --force "s3://$bucket"
}

function destroy_dynamo_table() {
    table="$(grep 'dynamodb_table' backend.tfvars | sed -e 's/dynamodb_table = //' -e 's/"//g')"
    aws dynamodb delete-table --table-name="$table"
}

function destroy_generated_files() {
    rm -f ./backend.tfvars
    rm -f ./main.tfvars
}

# Tear down all provisioned infra
function destroy() {
    # shellcheck disable=SC2086
    terraform plan -destroy -var-file=main.tfvars -out plans/destroy.planfile main
    read -r -p "Are you sure you want to run this plan? (y/n) "
    if [[ $REPLY =~ ^[yY]$ ]]; then
        terraform apply plans/destroy.planfile
        rm -f plans/destroy.planfile
    else
        exit 0
    fi
    read -r -p "Do you wish to destroy the Terraform state? (y/n) "
    if [[ $REPLY =~ ^[yY]$ ]]; then
        destroy_bucket
        destroy_dynamo_table
        rm -rf terraform.tfstate.d
        rm -rf .terraform
    else
        exit 0
    fi
    read -r -p "Do you want to delete the generated config files? (y/n) "
    if [[ $REPLY =~ ^[yY]$ ]]; then
        destroy_generated_files
    fi
    success "All generated infrastructure successfully removed!"
}

# Provision infrastructure
function provision() {
    # If INFRA_PREFIX has not been set yet, request it from user
    if [ -z "$INFRA_PREFIX" ]; then
        DEFAULT_INFRA_PREFIX=$(LC_ALL=C tr -dc 'a-z0-9' < /dev/urandom | fold -w 5 | head -n 1)

        warnb << EOF
# Infrastructure Prefix

In order to ensure that provisioned resources are unique, this script uses a 
unique prefix for all resource names and ids.

By default, a random 5 character alphanumeric string is generated for you, but 
if you wish to provide your own, now is your chance. This value will be stored 
in 'main.tfvars' so that you only need provide it once, but make sure you source 
control the file.
EOF

        read -r -p "What prefix should be used? (default is $DEFAULT_INFRA_PREFIX): "
        INFRA_PREFIX="$REPLY"
        if [ -z "$INFRA_PREFIX" ]; then
            INFRA_PREFIX="$DEFAULT_INFRA_PREFIX"
        fi
    fi

    if ! echo "$INFRA_PREFIX" | grep -E '^[a-z0-9]{3,5}$'; then
        errorb << EOF
The prefix '$INFRA_PREFIX' is invalid!

It must consist only of the lowercase characters a-z and digits 0-9,
and must be between 3 and 5 characters long.
EOF
    fi

    # EC2 key pairs
    if [ -z "$KEY_PAIR" ]; then
        KEY_PAIR="$(get_config 'key_name')"
        if [ -z "$KEY_PAIR" ]; then
            read -r -p "Please provide the name of the key pair to use with EC2 hosts: "
            KEY_PAIR="$REPLY"
            if [ -z "$KEY_PAIR" ]; then
                error "You must provide a valid key pair name!"
                exit 2
            fi
        fi
    fi

    if ! aws ec2 describe-key-pairs --key-names="$KEY_PAIR" 2>/dev/null; then
        if [ "$DRY_RUN" == "true" ]; then
            action "DRY RUN: Would have created an EC2 key pair"
        else
            info "The key pair '$KEY_PAIR' does not exist, creating..."
            if ! output=$(aws ec2 create-key-pair --key-name="$KEY_PAIR"); then
                error "$output\\nFailed to generate key pair!"
            fi
            echo "$output" | jq '.KeyMaterial' --raw-output > "$KEY_PAIR.privkey"
            success "Created keypair successfully! Private key has been saved to ./$KEY_PAIR.privkey"
        fi
    fi

    if [ -z "$SECRET_KEY_BASE" ]; then
        SECRET_KEY_BASE="$(get_config 'secret_key_base')"
        if [ -z "$SECRET_KEY_BASE" ]; then
            SECRET_KEY_BASE="$(openssl rand -base64 64 | tr -d '\n')"
        fi
    fi

    # Save variables used by Terraform modules
    if [ ! -f ./backend.tfvars ] && [ ! -f ./main.tfvars ]; then
        # shellcheck disable=SC2154
        region="$TF_VAR_region"
        if [ -z "$region" ]; then
            # Try to pull region from local config
            if [ -f "$HOME/.aws/config" ]; then
                if [ "$AWS_PROFILE" == "default" ]; then
                    region=$(awk '/\[default\]/{a=1;next}; /\[/{a=0}a' ~/.aws/config | grep 'region' | sed -e 's/region = //')
                else
                    #shellcheck disable=SC1117
                    region=$(awk "/\[profile $AWS_PROFILE\]/{a=1;next}; /\[/{a=0}a" ~/.aws/config | grep 'region' | sed -e 's/region = //')
                fi
            fi
        fi
        if [ -z "$region" ]; then
            read -r -p "What region should infrastructure be created in (us-east-2): "
            if [ -z "$REPLY" ]; then
                region='us-east-2'
            else
                region="$REPLY"
            fi
        fi
        bucket="$(get_config 'bucket')"
        if [ -z "$bucket" ]; then
            bucket="poa-terraform-state"
        fi
        dynamo_table="$(get_config 'dynamodb_table')"
        if [ -z "$dynamo_table" ]; then
            dynamo_table="poa-terraform-locks"
        fi
        # Backend config only!
        {
            echo "region = \"$region\""
            echo "bucket = \"${INFRA_PREFIX}-$bucket\""
            echo "dynamodb_table = \"${INFRA_PREFIX}-$dynamo_table\""
            echo "key = \"terraform.tfstate\""
        } > ./backend.tfvars
        # Other configuration needs to go in main.tfvars or init will break
        {
            echo "region = \"$region\""
            echo "bucket = \"$bucket\""
            echo "dynamodb_table = \"$dynamo_table\""
            echo "key_name = \"$KEY_PAIR\""
            echo "prefix = \"$INFRA_PREFIX\""
            echo "secret_key_base = \"$SECRET_KEY_BASE\""
        } > ./main.tfvars
    fi

    # No Terraform state yet, so this is a fresh run
    if [ ! -d .terraform ]; then
        terraform workspace new base setup
        terraform workspace select base setup
        # shellcheck disable=SC2086
        terraform init -backend-config=backend.tfvars setup
        # shellcheck disable=SC2086
        terraform plan -out plans/setup.planfile setup

        if [ "$DRY_RUN" == "false" ]; then
            # No need to show the plan, it has already been displayed
            SKIP_SETUP_PLAN="true"
        fi
    fi

    workspace="$(terraform workspace show)"

    # Setup hasn't completed yet, perhaps due to a dry run
    if [ -f plans/setup.planfile ]; then
        if [ -z "$SKIP_SETUP_PLAN" ]; then
            # Regenerate setup plan if not fresh
            # shellcheck disable=SC2086
            terraform plan -out plans/setup.planfile setup
        fi
        # Wait for user approval if we're going to proceed
        if [ "$SKIP_APPROVAL" == "false" ]; then
            read -r -p "Take a moment to review the generated plan, and press ENTER to continue"
        fi
        if [ "$DRY_RUN" == "true" ]; then
            action "DRY RUN: Would have executed Terraform plan for S3 backend as just shown"
            warn "Unable to dry run further steps until S3 backend has been created!"
            exit 0
        fi
        terraform apply plans/setup.planfile
        rm plans/setup.planfile
        # Migrate state to S3
        # shellcheck disable=SC2086
        terraform init -force-copy -backend-config=backend.tfvars base
    fi

    if [ "$workspace" == "base" ]; then
        # Switch to main workspace
        terraform workspace new main main
        terraform workspace select main main
    fi
    # shellcheck disable=SC2086
    terraform init -backend-config=backend.tfvars -var-file=main.tfvars main
    # Generate the plan for the remaining infra
    # shellcheck disable=SC2086
    terraform plan -var-file=main.tfvars -out plans/main.planfile main
    if [ "$SKIP_APPROVAL" == "false" ]; then
        read -r -p "Take a moment to review the generated plan, and press ENTER to continue"
    fi
    if [ "$DRY_RUN" == "true" ]; then
        action "DRY RUN: Would have executed the Terraform plan just shown"
    fi

    # Apply the plan to provision the remaining infra
    terraform apply plans/main.planfile
    rm plans/main.planfile
    success "Infrastructure has been successfully provisioned!"
}

# Print all resource ARNs tagged with prefix=INFRA_PREFIX
function resources() {
    if [ -z "$INFRA_PREFIX" ]; then
        error "No prefix set, unable to locate tagged resources"
        exit 1
    fi
    # Yes, stagging, blame Amazon
    aws resourcegroupstaggingapi get-resources \
        --no-paginate \
        --tag-filters="Key=prefix,Values=$INFRA_PREFIX" | \
    jq '.ResourceTagMappingList[].ResourceARN' --raw-output
}

# Provide test data for validation
function precheck() {
     # Save variables used by Terraform modules
    if [ ! -f ./ignore.tfvars ]; then
        {
            echo "bucket = \"poa-terraform-state\""
            echo "dynamodb_table = \"poa-terraform-locks\""
            echo "key = \"terraform.tfstate\""
            echo "key_name = \"poa\""
            echo "prefix = \"prefix\""
        } > ./ignore.tfvars
    fi
}

# Parse options for this script
VERBOSE=false
HELP=false
DRY_RUN=false

# Environment variables for Terraform
AWS_PROFILE="${AWS_PROFILE:-default}"

COMMAND=
while [ "$1" != "" ]; do
    param=$(echo "$1" | sed -re 's/^([^=]*)=/\1/')
    val=$(echo "$1" | sed -re 's/^([^=]*)=//')
    case $param in
        -h | --help)
            HELP=true
            ;;
        -v | --verbose)
            VERBOSE=true
            ;;
        --dry-run)
            DRY_RUN=true
            ;;
        --no-color)
            disable_color
            ;;
        --profile)
            AWS_PROFILE="$val"
            ;;
        --skip-approval)
            SKIP_APPROVAL="true"
            ;;
        --)
            shift
            break
            ;;
        *)
            COMMAND="$param"
            shift
            break
            ;;
    esac
    shift
done

# Turn on debug mode if --verbose was set
if [ "$VERBOSE" == "true" ]; then
    set -x
fi

# Set working directory to the project root
cd "$(dirname "${BASH_SOURCE[0]}")/.."

# Export AWS_PROFILE if a non-default profile was chosen
if [ ! "$AWS_PROFILE" == "default" ]; then
    export AWS_PROFILE
fi

# If cached prefix is in PREFIX file, then use it
if [ -z "$INFRA_PREFIX" ]; then
    if ls ./*.tfvars >/dev/null; then
        INFRA_PREFIX="$(get_config 'prefix')"
    fi
fi

# Override command if --help or -h was passed
if [ "$HELP" == "true" ]; then
    # If we ever want to show help for a specific command we'll need this
    # HELP_COMMAND="$COMMAND"
    COMMAND=help
fi

check_prereqs

case $COMMAND in
    help)
        help
        ;;
    provision)
        provision
        ;;
    destroy)
        destroy
        ;;
    resources)
        resources
        ;;
    precheck)
        precheck
        ;;
    destroy_setup)
        destroy_bucket
        destroy_dynamo_table
        ;;
    *)
        error "Unknown task '$COMMAND'. Try 'help' to see valid tasks"
        exit 1
esac                

exit 0