#!/usr/bin/env bash

set -eo pipefail
set -x

# default variables
: "${PORT:=8000}"
: "${SLEEP:=1}"
: "${TRIES:=60}"

usage() {
  echo "usage: bin/run flower|web|worker|scheduler"
  exit 1
}

wait_for() {
  tries=0
  echo "Waiting for $1 to listen on $2..."
  while true; do
    [[ $tries -lt $TRIES ]] || return
    (echo > /dev/tcp/$1/$2) >/dev/null 2>&1
    result=
    [[ $? -eq 0 ]] && return
    sleep $SLEEP
    tries=$((tries + 1))
  done
}

gcp_default_extras() {
python - <<END
import json

extra = {
    "extra__google_cloud_platform__project": "dummy_project_id",
    "extra__google_cloud_platform__keyfile_dict": json.dumps({
        "type": "service_account",
        "project_id": "dummy-project",
        "private_key_id": "d8b87911500d92da857b92ab74a9d53631e8904e",
        "private_key": "-----BEGIN PRIVATE KEY-----\nMIIEogIBAAKCAQEAjj1EQOX7o14lnDRJAyFZeXGNLywPuyf/vdDt6tI2v0cTgoOH\nEAaCL5OYyXnYnLYEfQIap5sU8DcVDgiidCQv3+PKyzXPV8YPe9E2ureMZJAu48Zn\nKPzoltvLFFqAIX8yzezb3Av+tnRY495iuqTpR1x4ZKXvgVPFtOotT1dseZhNbZ2x\nmg5mNVxxsV3pVPmaRLj8CjtJI2qU3EZ2AbxTON3k4k0Q+7Gy2bG5Ef5rKbTFsp5O\nFF7MGjX7A54H7keVCc0QkEWSx/bsaRkGsnDbrZG3Kn118xJOYY/+CREURNO1FNTA\nzlJgIxN4PAOAGRXEBeXcsCZKp3MYpU75cvF5/wIDAQABAoIBAFtXkMs0ZaKFxRVI\nplJySiko+IeAfiGsEBlvYDnaAPpYxHidylBKPbQbzpQjwSzx3nQAs+lKN+oDFWxL\nszduPahDemmBBsPRFwRmWAUT9f5mcRYoxPqXXy3xu7o4W+wm6RNAtffbZBj7IlJC\n75f4ay4+fbn0rZeZmm8Rq0M2WxzB8DDiO96N3JLhPYV/uR9TeQp4pn/lUYLKRO2i\nuMBmCLl6hjaV8/0Dox0KhNICox+BXKnIUlIrDV9NanmO5dAZJPOzSHH6aJ4ovdFo\nmLdqlBw50qGm9DaYHR3pbZo17kX/42yZogrcEJMNGHjJgZna+W54gVsP6Gs4ht3G\nO+O6A4ECgYEA2oASmeAVs65VdVpBpI6Ky1PW4wFV+zMxVykbgY++6ZTg8rjRfzps\nDsmoGTaI3b/lxeiakIu96bkSr3A7r1fVuKtwSzU5jSSkONe8mWiHt4VIX+x5GiY2\nP4CZCEi3v+bVzovmXcx9Wy62h+bBWsrkYav2CrXrOSB7Byz7XduQVmECgYEApqag\nrpY7Db2HuCK3yhp08AT7itDi9ck43Zes/YB2TqXtTYsDVWVNeYcXkSh1VrQytQZW\ncPXtfmBXogGSTnjXEmpCRf+z5Bsxu57Zole5ENiBTzEcpFJpQ/5UPH8PJzWA+dB4\ncAuvgRKxnTG8LRhkrbVM8t0WJi1v1nP6DyDA7F8CgYARC/bnjGUFDK/cJPuEFB7d\n+B+GvF7x5y+NRkbAF+/kF1ppdWPa0jsF+FOmC+wnqMYLZ7dPWaeqaWb8yvvNFUQ2\nUSHErFVeHqK1UJeFPHOCLOLVoQRdtud6ktTdoZa8YQ0DPUTuwnpxN7bD6YviQnwI\n5rqeYU0FuvP+PlMqImwjYQKBgG/fbgCls1D/CcwP6mdPKW8zORWwMpwjD/yZ5LRs\n937Gnq4ugvdhwQezK5vzmAmzgFLLxV2himQLEukbuvbY4jBnisPo6v9XTiSQd7Mm\nxoLLhMr/wiWBeU7+vde6yBZfMY0CaMd24MN6JCfNinCPbo66JcTnrAXG/MvvIU/k\ngf5rAoGAfFN3JUFRcGZsQ+2qqeghcnD38iXkhBE5WL/g5CMA+v/AXvS7L67HVcTa\nV2O/KlmS53RV5MPnNZsOCyWBtohLYc+8kTAwl+DUf2ZfbzLJty6mCPrlaRaqCZ9q\nXTu2Ea7bs7gbVmhj2xcD+OPBUkITy7+XWAqwqe6C4hhwiHTDIoI=\n-----END PRIVATE KEY-----\n",
        "client_email": "test-account@dummy-project.iam.gserviceaccount.com",
        "client_id": "19327456837",
        "auth_uri": "https://accounts.google.com/o/oauth2/auth",
        "token_uri": "https://oauth2.googleapis.com/token",
        "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
        "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/test-account%40dummy-project.iam.gserviceaccount.com"
    })
}

print(json.dumps(extra))
END
}

aws_default_extras() {
python - <<END
import json

extra = {
    "aws_access_key_id": "dummy_access_key_id",
    "aws_secret_access_key": "dummy_secret_access_key",
}

print(json.dumps(extra))
END
}

init_connections() {
    # set dummy credentials
    export AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-dummy_access_key_id}
    export AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-dummy_secret_access_key}

    airflow connections --delete --conn_id databricks_default

    airflow connections --add \
        --conn_id databricks_default \
        --conn_type databricks \
        --conn_host https://dbc-caf9527b-e073.cloud.databricks.com \
        --conn_extra "{\"token\":\"${DB_TOKEN}\", \"host\": \"\"}"

    gcp_conn=(
      "google_cloud_airflow_dataproc"
      "google_cloud_derived_datasets"
      "google_cloud_prio_a"
      "google_cloud_prio_admin"
      "google_cloud_prio_b"
      "google_cloud_glam_fenix_dev"
    )
    for conn_id in "${gcp_conn[@]}"; do
        airflow connections --delete --conn_id "${conn_id}"
        airflow connections --add \
        --conn_id "${conn_id}" \
        --conn_type google_cloud_platform \
        --conn_extra "$(gcp_default_extras)"
    done

    aws_conn=(
      "airflow_taar_rw_s3"
      "amplitude_s3_conn"
      "aws_dev_iam_s3"
      "aws_data_iam_s3"
      "aws_dev_socorro_telemetry_parquet_s3"
      "aws_dev_telemetry_public_analysis_2_rw"
      "aws_prod_fx_usage_report"
      "aws_prod_probe_scraper"
      "aws_socorro_readonly_s3"
    )
    for conn_id in "${aws_conn[@]}"; do
        airflow connections --delete --conn_id "${conn_id}"
        airflow connections --add \
        --conn_id "${conn_id}" \
        --conn_type s3 \
        --conn_extra "$(aws_default_extras)"
    done
}

init_variables() {
    airflow variables -s "bugzilla_probe_expiry_bot_api_key" "bugzilla-api-key"
    airflow variables -s "app_store_connect_username" "username"
    airflow variables -s "app_store_connect_password" "password"
    airflow variables -s "LEANPLUM_FIREFOX_ANDROID_BETA_APP_ID" "firefox-nightly-app-id"
    airflow variables -s "LEANPLUM_FIREFOX_ANDROID_BETA_CLIENT_KEY" "firefox-nightly-client-key"
    airflow variables -s "LEANPLUM_FIREFOX_ANDROID_NIGHTLY_APP_ID" "firefox-beta-app-id"
    airflow variables -s "LEANPLUM_FIREFOX_ANDROID_NIGHTLY_CLIENT_KEY" "firefox-beta-client-key"
    airflow variables -s "LEANPLUM_FIREFOX_ANDROID_RELEASE_APP_ID" "firefox-release-app-id"
    airflow variables -s "LEANPLUM_FIREFOX_ANDROID_RELEASE_CLIENT_KEY" "firefox-release-client-key"
    airflow variables -s "LEANPLUM_FIREFOX_IOS_RELEASE_APP_ID" "firefox-ios-app-id"
    airflow variables -s "LEANPLUM_FIREFOX_IOS_RELEASE_CLIENT_KEY" "firefox-ios-client-key"
    airflow variables -s "surveygizmo_daily_attitudes_survey_id" "12345"
    airflow variables -s "surveygizmo_api_token" "tokentokentoken"
    airflow variables -s "surveygizmo_api_secret" "tapsekret"
    airflow variables -s "jetstream_cluster_ip" "127.0.0.1"
    airflow variables -s "jetstream_cluster_cert" "cert"

    airflow variables -s "taar_bigtable_instance_id" "taar_bigtable_instance_id"
    airflow variables -s "taar_etl_storage_bucket" "taar_etl_storage_bucket"
    airflow variables -s "taar_gcp_project_id" "taar_gcp_project_id"
    airflow variables -s "taar_dataflow_subnetwork" "taar_dataflow_subnetwork"
}

[ $# -lt 1 ] && usage

# Only wait for backend services in development
# http://stackoverflow.com/a/13864829
[ ! -z ${DEVELOPMENT+check} ] && wait_for db 3306 && wait_for redis 6379

case $1 in
  flower)
    exec airflow flower
    ;;
  web)
    airflow initdb
    airflow upgradedb

    # Only init connections in dev
    [ ! -z ${DEVELOPMENT+check} ] && init_connections && init_variables
    echo "[testing_stage_0]: Initialization complete!"

    exec airflow webserver -p ${PORT} --workers 4
    ;;
  worker)
    exec airflow worker
    ;;
  scheduler)
    exec airflow scheduler
    ;;
  *)
    exec "$@"
    ;;
esac