#!/usr/bin/env bash
set -e

# Default Values
INSTANCE_TYPE="c5.xlarge"

MONIKER=

GREEN='\033[0;32m'
RED='\033[0;31m'
NO_COLOR='\033[0m'

check_aws_connectivity() {
    echo "Checking aws connectivity..."
    USERNAME=$(aws iam get-user --query "User.UserName" --output text)
    if [ "$USERNAME" == None ]; then
        echo "Failed to get valid aws username."
        exit 1
    fi
}

check_aws_connectivity

confirm_create() {
    if [ -z "$1" ]; then
        echo -e "Create dock with the following values?"
        echo -e "Instance type: ${GREEN}${INSTANCE_TYPE}${NO_COLOR}"
        echo -e "Name: ${GREEN}${INSTANCE_NAME}${NO_COLOR}"
        echo -e "AMI ID: ${GREEN}${AMI_ID}${NO_COLOR}"
        read -e -p "Type enter to Cancel, h for Help, y to Create: " RESPONSE
    fi

    if [ "$RESPONSE" == "h" ]; then print_help; fi
}

print_help() {
    echo "Create dock - Help"
    echo
    echo "Description"
    echo "  This script uses the aws cli to create a new ec2 dock instance from the latest AMI."
    echo "  If no options are passed into the script, it will prompt with defaults."
    echo "  The register-dock script is run automatically after the instance is ready."
    echo
    echo "  A number of values used in creating the ec2 instance (security-group-ids, etc.) can be over-ridden"
    echo "  by specifying a ~/.docker/dock.cfg file with the following values:"
    echo
    echo '    key_name="..."'
    echo '    security_group_ids="..."'
    echo '    iam_instance_profile="..."'
    echo '    subnet_id="..."'
    echo '    tag_specifications="..."'
    echo
    echo "Usage"
    echo "  $ create-dock [options]"
    echo
    echo "Options"
    echo -e "\n -n dock-name\n    Specify a name for this dock with this format: my-dock-name"
    echo -e "    If dock-name is not specified, the name defaults to ${GREEN}${INSTANCE_NAME}${NO_COLOR}"
    echo -e "\n -m moniker\n    moniker will be added to your cli prompt to indicate that you are docked. (Default: instance IP address)"
    echo
    echo "    Example CLI prompt with python virtual environment enabled and docked to remote worker:"
    echo -e "    $ (venv-name) prompt$ [dock:moniker] "

    echo -e "\n -i instance-type\n    The default instance type is ${INSTANCE_TYPE} (no GPU).\n    Other options include p2.xlarge (GPU), m5.2xlarge, etc."
    echo -e "\n -a ami-id\n    (Optional) The ami to use for the instance."
    echo -e "\n -s skip-regisration\n    (Optional) Skip running register-dock."
    echo -e "\n -h help\n    This help"
    echo
    echo "Examples"
    echo
    echo -e "  $ create-dock\n    Create default instance type ($INSTANCE_TYPE) named ${INSTANCE_NAME}.\n    CLI prompt moniker when docked will be [dock:IP address]."
    echo
    echo -e "  $ create-dock -i t2.micro -m my-dock\n    Create t2.micro instance type named ${INSTANCE_NAME}.\n    CLI prompt moniker when docked will be [dock:my-dock]."


    exit 0
}

get_private_ip() {
    aws ec2 describe-instances \
        --filters Name=instance-id,Values="$1" \
        --query 'Reservations[*].Instances[*].PrivateIpAddress' --output text
}


# Parse command line arguments in any order
hflag=''
nflag=''    # dock name flag
iflag=''    # instance type
aflag=''    # ami flag
mflag=''    # moniker flag
regflag='true'
while getopts 'hn:i:a:m:s' flag; do    # if a character is followed by a colon, that argument is expected to have an argument.
  case "${flag}" in
    h) hflag='true';;
    n) nflag='true'; INSTANCE_NAME="${OPTARG}" ;;
    i) iflag='true'; INSTANCE_TYPE="${OPTARG}" ;;
    a) aflag='true'; AMI_ID="${OPTARG}" ;;
    m) mflag='true'; MONIKER="${OPTARG}" ;;
    s) regflag='';;
    *) error "Unexpected option ${flag}" ;;
  esac
done

INSTANCE_NAME="${USERNAME}-${MONIKER:-dock}"

# Help
if [ -n "$hflag" ] || [ "$RESPONSE" == "h" ]; then
    print_help
fi

if [ -z "$AMI_ID" ]; then
    AMI_ID=$(aws ec2 describe-images --filters "Name=tag:Name,Values=resero-labs-dlami" --query 'Images[*].{ID:ImageId}' --output text)
    if [ -z "$AMI_ID" ]; then
        echo "Unable to fetch default AMI ID."
        exit 1
    fi
fi

# Confirmation
confirm_create
if [ "$RESPONSE" != "y" ] && [ "$RESPONSE" != "h" ]; then
    echo "Canceled"
    exit 0
fi


# Create dock
if [ -e "$HOME/.docker/dock.cfg" ]; then
    source "$HOME/.docker/dock.cfg"
else
    key_name="resero-staging"
    security_group_ids="sg-213eb35a sg-3bde0341 sg-b93e0dc2 sg-1bd90461"
    iam_instance_profile="lanista-app"
    subnet_id="subnet-b8b440de"
    tag_specifications="ResourceType=instance,Tags=[{Key=Name,Value=${INSTANCE_NAME}},\
            {Key=business_unit,Value='Archiving & Governance'},{Key=component,Value='ec2 instance'},\
            {Key=product,Value='Resero Development'},{Key=support_level,Value=dev},\
            {Key=created_by,Value=${USERNAME}}]"
fi
if [ ! -f "$HOME/.ssh/$key_name" ]; then
    echo
    echo
    echo -e "${RED}It appears the the key required to access the EC2 instance doesn't exist ($HOME/.ssh/$key_name).${NO_COLOR}"
    echo
    echo "Please ensure that the correct ssh-key is configured in ~/.docker/dock.cfg"
    exit 1
fi
INSTANCE_ID=$(aws ec2 run-instances \
    --subnet-id ${subnet_id} \
    --image-id $AMI_ID \
    --instance-type "${INSTANCE_TYPE}" \
    --disable-api-termination \
    --block-device-mappings "DeviceName='/dev/sda1',Ebs={VolumeSize=100,VolumeType='gp2'}" \
    --tag-specifications "${tag_specifications}" \
    --iam-instance-profile Name="${iam_instance_profile}" \
    --key-name "${key_name}" \
    --security-group-ids ${security_group_ids} \
    | grep InstanceId | awk -F '"' '{print $4}' \
    )

if [ ! -z "$INSTANCE_ID" ]; then
    echo "Creating instance id: $INSTANCE_ID"
    echo "Waiting for instance to start..."
    aws ec2 wait system-status-ok --instance-ids $INSTANCE_ID

    IP_ADDRESS=$(get_private_ip $INSTANCE_ID)

    if [ -n "$regflag" ]; then
        if [ $(echo $IP_ADDRESS | grep -c -E '[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}') == 1 ]; then
            echo 'Registering secure remote docker api'
            register-dock -f ubuntu "$IP_ADDRESS" "$MONIKER" "$INSTANCE_ID"
        fi
    else
        mkdir -p ~/.docker/${IP_ADDRESS}
        printf "DOCK_USER=ubuntu\nDOCK_MONIKER=$MONIKER\nDOCK_HOSTNAME=$IP_ADDRESS\nDOCK_IP=$IP_ADDRESS\n" > $HOME/.docker/${IP_ADDRESS}/connection_config.txt
    fi
    printf "DOCK_INSTANCE_ID=$INSTANCE_ID" >> $HOME/.docker/${IP_ADDRESS}/connection_config.txt &&


    echo "New EC2 instance available to dock at $IP_ADDRESS"
    echo "try these commands:"
    echo " $ source dock $IP_ADDRESS"
    echo " $ castoff"
    echo " $ ssh-add -K ~/.ssh/private-key-name"
    echo " $ ssh ubuntu@$IP_ADDRESS"
else
    echo "Failed to create instance."
fi
