#!/usr/bin/env bash
# ORIGINAL SOURCE: https://github.com/nicor88/aws-ecs-airflow/


TRY_LOOP="20"

: "${REDIS_HOST:="redis"}"
: "${REDIS_PORT:="6379"}"

: "${POSTGRES_HOST:="postgres"}"
: "${POSTGRES_PORT:="5432"}"

: "${AIRFLOW_UI_PASS_ARN:="arn:aws:secretsmanager:eu-west-1:776891437216:secret:airflow-ui-pass-Z8rpW1"}"
: "${BITBUCKET_SSH_KEY_ARN:="arn:aws:secretsmanager:eu-west-1:776891437216:secret:bitbucket-fansifter-airflow-access-key-Ll5zJj"}"
: "${BITBUCKET_URL:="git@bitbucket.org:fansifter/fansifter-airflow.git"}"

wait_for_port() {
  local name="$1" host="$2" port="$3"
  local j=0
  while ! nc -z "$host" "$port" >/dev/null 2>&1 < /dev/null; do
    j=$((j+1))
    if [ $j -ge $TRY_LOOP ]; then
      echo >&2 "$(date) - $host:$port still not reachable, giving up"
      exit 1
    fi
    echo "$(date) - waiting for $name... $j/$TRY_LOOP"
    sleep 5
  done
}

get_ssh_keys(){
  local credentials
  credentials="$(aws secretsmanager get-secret-value --secret-id $BITBUCKET_SSH_KEY_ARN --query SecretString --output text)"

  mkdir -p /usr/local/airflow/.ssh
  echo "$(echo $credentials | jq -r '.public_key' | base64 --decode)" \
    > /usr/local/airflow/.ssh/id_rsa.pub
  echo "$(echo $credentials | jq -r '.private_key' | base64 --decode)" \
    > /usr/local/airflow/.ssh/id_rsa
  chmod 600 /usr/local/airflow/.ssh/id_rsa
  chmod 600 /usr/local/airflow/.ssh/id_rsa.pub
  chmod 700 /usr/local/airflow/.ssh
  eval "$(ssh-agent -s)"
  ssh-add /usr/local/airflow/.ssh/id_rsa
}

pull_dags_and_plugins(){
  if [ $IS_LOCAL = "true" ]; then
     echo "Running Airflow locally, no need to pull DAGs and plugins from remote..."
  else
    echo "Pulling DAGs and plugins from remote Bitbucket repo..."
    export GIT_SSH_COMMAND="ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no"

    case $STAGE in
      dev)
        BRANCH="master"
        ;;
      test)
        BRANCH="deploy_test"
        ;;
      live)
        BRANCH="deploy_live"
        ;;
      *)
        exit 1
        ;;
    esac

    mkdir -p /mnt/efs/fansifter-airflow
    cd /mnt/efs/fansifter-airflow
    
    git init > /dev/null 2>&1
    git remote add -f origin $BITBUCKET_URL > /dev/null 2>&1
    git fetch origin $BRANCH
    git checkout origin/$BRANCH -- dags
    git checkout origin/$BRANCH -- plugins
  fi
}


wait_for_port "Postgres" "$POSTGRES_HOST" "$POSTGRES_PORT"

wait_for_port "Redis" "$REDIS_HOST" "$REDIS_PORT"

case "$1" in
  webserver)
    airflow db init
		sleep 5
		airflow users create --username fansifter \
		                     --password $(aws secretsmanager get-secret-value --secret-id $AIRFLOW_UI_PASS_ARN --query SecretString --output text | jq -r '.password') \
		                     --firstname Fan \
		                     --lastname Sifter \
		                     --role Admin \
		                     --email airflow@fansifter.com > /dev/null 2>&1
		sleep 5
    exec airflow webserver
    ;;
  worker)
    sleep 10
    get_ssh_keys
    exec airflow celery worker
    ;;
  scheduler)
    sleep 10
    get_ssh_keys
    pull_dags_and_plugins
    exec airflow scheduler
    ;;
  flower)
    sleep 10
    exec airflow celery flower
    ;;
  version)
    exec airflow "$@"
    ;;
  *)
    exec "$@"
    ;;
esac