#Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#Permission is hereby granted, free of charge, to any person obtaining a copy of
#this software and associated documentation files (the "Software"), to deal in
#the Software without restriction, including without limitation the rights to
#use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
#the Software, and to permit persons to whom the Software is furnished to do so.
#THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
#FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
#COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
#IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
#CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

if [ $# -ne 2 ]
  then
    echo "Usage: $0 <EKS NameSpace> <EMR Virtual Cluster Name>"
    exit
fi

region=`cat /tmp/region`
NameSpace=$1
Virtual_Cluster_Name=$2

stackName=`sudo cat /tmp/stackName | sed 's/-PrepStack-.*//g'`
kubectl create ns $NameSpace

EKS_Cluster_Name=`sudo cat /tmp/clusterName`

#Enable cluster access for Amazon EMR on EKS
eksctl create iamidentitymapping \
    --cluster $EKS_Cluster_Name \
    --namespace $NameSpace  \
    --service-name "emr-containers" \
    --region $region

#Enable IAM Roles for Service Accounts (IRSA) on the EKS cluster
/usr/local/bin/aws eks describe-cluster --name $EKS_Cluster_Name --query "cluster.identity.oidc.issuer" --region $region --output text

#To associate of above IAM OIDC identity provider with your cluster with eksctl
eksctl utils associate-iam-oidc-provider --cluster $EKS_Cluster_Name --approve --region $region

#Create trust-execution-role, adding the oidc identity provider as principal
#to the trust policy on the execution role
EMR_Execution_Role=`/usr/local/bin/aws cloudformation describe-stacks --stack-name $stackName --region $region | jq -r  .Stacks[].Outputs[] |  grep -w -A1 EMRExecutionRole | tail -1 | awk -F\" '{print $4}'`

/usr/local/bin/aws emr-containers update-role-trust-policy \
       --cluster-name $EKS_Cluster_Name \
       --namespace $NameSpace \
       --role-name $EMR_Execution_Role \
       --region $region 

#Register virtual cluster with EKS
cat << EOF > /tmp/registering-eks-cluster
/usr/local/bin/aws emr-containers create-virtual-cluster \
--name $Virtual_Cluster_Name \
--container-provider '{
    "id": "$EKS_Cluster_Name",
    "type": "EKS",
    "info": {
        "eksInfo": {
            "namespace": "$NameSpace"
        }
} }' \
--region $region 
EOF

sh /tmp/registering-eks-cluster

echo
echo
echo

#Preparing the dag for MWAA submission
airflow_bucket=`cat /tmp/bucket`
sleep 2
virtual_cluster_id=`/usr/local/bin/aws emr-containers list-virtual-clusters --region $region --output text | grep -i -w RUNNING | grep -w $Virtual_Cluster_Name | awk '{print $4}'`
account_id=`aws sts get-caller-identity --region $region --output text | awk '{print $1}'`
emr_execution_role_arn=`/usr/local/bin/aws cloudformation describe-stacks --stack-name $stackName --region $region | jq -r  .Stacks[].Outputs[] |  grep -w -A1 EMRExecutionRoleArn | tail -1 | awk -F\" '{print $4}'`

sudo sed -e s#AIRFLOW_BUCKET#$airflow_bucket# -e s#VIRTUAL_CLUSTER_ID#$virtual_cluster_id# -e s#EMR_EXECUTION_ROLE_ARN#$emr_execution_role_arn# /tmp/citibike_all_dag.py.template > /tmp/citibike_all_dag.py

/usr/local/bin/aws s3 cp /tmp/citibike_all_dag.py s3://$airflow_bucket/dags/ --quiet
echo
echo "EMR Virtual cluster: $Virtual_Cluster_Name in EC2 namespace: $NameSpace has been created" 
echo 
echo "DAG file /tmp/citibike_all_dag.py has been updated with the current S3 bucket, virtual cluster id and account id and is copied to s3://$airflow_bucket/dags/"

mwaaenv=`/usr/local/bin/aws cloudformation describe-stacks --stack-name $stackName --region $region | jq -r  .Stacks[].Outputs[] | grep -w -A1 MWAAEnvName | tail -1 | awk -F\" '{print $4}'`
host=`aws mwaa create-web-login-token --name $mwaaenv --region $region | jq -r .WebServerHostname`
token=`aws mwaa create-web-login-token --name $mwaaenv --region $region | jq -r .WebToken`
url="https://$region.console.aws.amazon.com/mwaa/home?region=$region#environments/"
echo
echo "Now point your browser at URL below, open the UI for environment $mwaaenv to see Airflow in action!"
echo
echo $url