#!/bin/bash

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#########################
# The command line help #
#########################
usage() {
    echo "Usage: $0"
    echo "   -h | --hadoop, hadoop-version"
    echo "   -s | --spark, spark version"
    echo "   -d | --directory, workspace directory"
    exit 1
}

case "$1" in
   --help)
       usage
       ;;
esac

while getopts ":h:s:d:-:" opt; do
  case $opt in
    h) SPARK_HADOOP_VERSION="$OPTARG"
    printf "Specified hadoop version is: %s\n" "$SPARK_HADOOP_VERSION"
    ;;
    s) SPARK_VERSION="$OPTARG"
    printf "Specified spark version is: %s\n" "$SPARK_VERSION"
    ;;
    d) WORK_SPACE_ROOT="$OPTARG"
    printf "Specified workspace directory is: %s\n" "$WORK_SPACE_ROOT"
    ;;
    -)
      case "$OPTARG" in
        hadoop)
          SPARK_HADOOP_VERSION="${!OPTIND}"; OPTIND=$(( OPTIND + 1 ))
          printf "Specified hadoop version is: %s\n" "$SPARK_HADOOP_VERSION"
          ;;
        spark)
          SPARK_VERSION="${!OPTIND}"; OPTIND=$(( OPTIND + 1 ))
          printf "Specified hadoop version is %s\n" "$SPARK_VERSION"
          ;;
        directory)
          WORK_SPACE_ROOT="${!OPTIND}"; OPTIND=$(( OPTIND + 1 ))
          printf "Specified workspace directory is: %s\n" "$WORK_SPACE_ROOT"
          ;;
        *) echo "Invalid option --$OPTARG" >&2
          ;;
    esac ;;
    \?) echo "Invalid option -$OPTARG" >&2
    ;;
  esac
done

if [ -z "$WORK_SPACE_ROOT" ]; then
  SCRIPT_PATH=$(cd "$(dirname "$0")" || pwd)
  WORK_SPACE_ROOT=$(dirname "$SCRIPT_PATH")
  echo "Using default workspace dir: $WORK_SPACE_ROOT"
fi

if [ -z "$SPARK_HADOOP_VERSION" ]; then
  SPARK_HADOOP_VERSION=3
  echo "Using default spark hadoop version: $SPARK_HADOOP_VERSION"
fi

if [ -z "$SPARK_VERSION" ]; then
  SPARK_VERSION=3.3.1
  echo "Using default spark version: $SPARK_VERSION"
fi

SPARK_BINARY_FILE_NAME=spark-${SPARK_VERSION}-bin-hadoop${SPARK_HADOOP_VERSION}.tgz
SPARK_BINARY_FILE_PATH=$WORK_SPACE_ROOT/$SPARK_BINARY_FILE_NAME
SPARK_BINARY_FILE_CHECKSUM_FILE_NAME=$SPARK_BINARY_FILE_NAME.sha512
SPARK_BINARY_FILE_CHECKSUM_FILE_PATH=$WORK_SPACE_ROOT/$SPARK_BINARY_FILE_CHECKSUM_FILE_NAME
FORMATTED_SPARK_BINARY_FILE_CHECKSUM_FILE_NAME=$SPARK_BINARY_FILE_CHECKSUM_FILE_NAME.formatted
FORMATTED_SPARK_BINARY_FILE_CHECKSUM_FILE_PATH=$WORK_SPACE_ROOT/$FORMATTED_SPARK_BINARY_FILE_CHECKSUM_FILE_NAME

SPARK_HOME=$WORK_SPACE_ROOT/spark-${SPARK_VERSION}-bin-hadoop${SPARK_HADOOP_VERSION}
SPARK_EXAMPLE_JAR=local:///opt/spark/examples/jars/spark-examples_2.12-${SPARK_VERSION}.jar

K8S_ENDPOINT=http://localhost:8001
SPARK_EXECUTOR_NUM=1
SPARK_DOCKER_IMAGE=docker.io/apache/spark:v${SPARK_VERSION}

if [ -f "$SPARK_BINARY_FILE_PATH" ]; then
  echo "The binary file $SPARK_BINARY_FILE_NAME has been cached!"
else
  echo "The binary file $SPARK_BINARY_FILE_NAME did not exist, try to download."
  wget "https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${SPARK_BINARY_FILE_NAME}" -O "${SPARK_BINARY_FILE_PATH}"
fi

if [ -f "$SPARK_BINARY_FILE_CHECKSUM_FILE_PATH" ]; then
  echo "The binary checksum file $SPARK_BINARY_FILE_CHECKSUM_FILE_NAME has been cached!"
else
  echo "The binary checksum file $SPARK_BINARY_FILE_CHECKSUM_FILE_NAME did not exist, try to download."
  wget "http://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/${SPARK_BINARY_FILE_CHECKSUM_FILE_NAME}" -O "${SPARK_BINARY_FILE_CHECKSUM_FILE_PATH}"
fi

if [[ "$SPARK_VERSION" < "3.2.1"  || "$SPARK_VERSION" == "3.2.1" ]]; then
  if [ -f "$FORMATTED_SPARK_BINARY_FILE_CHECKSUM_FILE_PATH" ]; then
    echo "The formatted binary checksum file $FORMATTED_SPARK_BINARY_FILE_CHECKSUM_FILE_NAME has been cached!"
  else
    # format the official checksum file for verifying
    echo "The formatted binary checksum file $FORMATTED_SPARK_BINARY_FILE_CHECKSUM_FILE_NAME did not exist, try to generate."
    < "$SPARK_BINARY_FILE_CHECKSUM_FILE_PATH" tr -d " \t\n\r" | awk -v awkvar="$SPARK_BINARY_FILE_PATH" -F: '{print $2 "  " awkvar}' > "$FORMATTED_SPARK_BINARY_FILE_CHECKSUM_FILE_PATH"
  fi
  # check signature to verify the completeness
  if [[ 'OK' == $(shasum -c -a 512 "$FORMATTED_SPARK_BINARY_FILE_CHECKSUM_FILE_PATH"  | awk '{print $2}') ]]; then
    echo "The checksum is matched!"
    echo "Try to remove the old unpacked dir and re-uncompress it"
    rm -rf "$WORK_SPACE_ROOT"/spark-"${SPARK_VERSION}"-bin-hadoop"${SPARK_HADOOP_VERSION}"
    tar -xvzf "$SPARK_BINARY_FILE_PATH" -C "$WORK_SPACE_ROOT"
  else
    echo "The checksum is not matched, Removing the incompleted file, please download it again."
    rm -f "$SPARK_BINARY_FILE_PATH"
    exit 0
  fi
else
  # check signature to verify the completeness
  if [[ $(shasum -a 512 "$SPARK_BINARY_FILE_PATH" | awk '{print $1}') == $(awk '{print $1}' < "$SPARK_BINARY_FILE_CHECKSUM_FILE_NAME") ]]; then
    echo "The checksum is matched!"
    echo "Try to remove the old unpacked dir and re-uncompress it"
    rm -rf "$WORK_SPACE_ROOT"/spark-"${SPARK_VERSION}"-bin-hadoop"${SPARK_HADOOP_VERSION}"
    tar -xvzf "$SPARK_BINARY_FILE_PATH" -C "$WORK_SPACE_ROOT"
  else
    echo "The checksum is not matched, Removing the incompleted file, please download it again."
    rm -f "$SPARK_BINARY_FILE_PATH"
    exit 0
  fi
fi

# Create a namespace
cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Namespace
metadata:
  name: spark-test
EOF

# Create service account and role bindings inside the spark-test namespace:
cat <<EOF | kubectl apply -n spark-test -f -
apiVersion: v1
kind: ServiceAccount
metadata:
  name: spark
  namespace: spark-test
---
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
  name: spark-role
  namespace: spark-test
rules:
- apiGroups: [""]
  resources: ["pods"]
  verbs: ["get", "watch", "list", "create", "delete"]
- apiGroups: [""]
  resources: ["configmaps"]
  verbs: ["get", "create", "delete"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
  name: spark-role-binding
  namespace: spark-test
subjects:
- kind: ServiceAccount
  name: spark
  namespace: spark-test
roleRef:
  kind: Role
  name: spark-role
  apiGroup: rbac.authorization.k8s.io
EOF

# spark submit command
"${SPARK_HOME}"/bin/spark-submit \
  --master k8s://${K8S_ENDPOINT} --deploy-mode cluster --name spark-pi \
  --class org.apache.spark.examples.SparkPi \
  --conf spark.executor.instances=${SPARK_EXECUTOR_NUM} \
  --conf spark.kubernetes.namespace=spark-test \
  --conf spark.kubernetes.executor.request.cores=1 \
  --conf spark.kubernetes.container.image="${SPARK_DOCKER_IMAGE}" \
  --conf spark.kubernetes.authenticate.driver.serviceAccountName=spark \
  --conf spark.kubernetes.driver.podTemplateFile=../driver.yaml \
  --conf spark.kubernetes.executor.podTemplateFile=../executor.yaml \
  "${SPARK_EXAMPLE_JAR}"
