#!/usr/bin/env 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.

if [ $# -lt 2 ] || [ "$1" == "-h" ] || [ "$1" == "--help" ]; then
  echo "Usage: ./upload-fsimage.sh image-txid remote-dir [ name-dir ]"
  echo "       Takes the steps necessary to place all necessary components for"
  echo "       the fsimage with the specified transaction ID (leading 0s not"
  echo "       required) onto HDFS. This includes the fsimage file itself,"
  echo "       the MD5 hash, the VERSION file, and the XML version of the"
  echo "       fsimage file, which will be temporarily generated by this script"
  echo "       (into \$TMPDIR if specified, else /tmp)."
  echo "       These files will be uploaded to remote_dir (created if it does"
  echo "       not yet exist). This is the same directory that should be passed"
  echo "       to the Client as fs_image_dir."
  echo "       If name-dir is specified, looks for fsimage files under"
  echo "       \${name-dir}/current. Otherwise, looks in the current directory."
  exit 1
fi
image_txid="$1"
remote_dir="$2"
if [[ $# -ge 3 ]]; then
  name_dir="$3/current"
else
  name_dir="$(pwd)"
fi

image_file_count="$(find -H "${name_dir}" -maxdepth 1 -mindepth 1 -name "fsimage_*$image_txid" -type f | wc -l)"
if [[ "$image_file_count" != 1 ]]; then
  echo "Error; found $image_file_count matching fsimage files."
  exit 1
fi
image_file="$(find -H "${name_dir}" -maxdepth 1 -mindepth 1 -name "fsimage_*$image_txid" -type f)"
image_file_name="$(basename "${image_file}")"
echo "Using fsimage: $image_file_name"
image_file_md5="${image_file}.md5"
version_file="${name_dir}/VERSION"

if ! [[ -f "$image_file_md5" ]]; then
  echo "Unable to find MD5 file for fsimage; looked at $image_file_md5"
  exit 1
fi
if ! [[ -f "$version_file" ]]; then
  echo "Unable to find VERSION file; looked at $version_file"
  exit 1
fi

if ! tmp_dir="$(mktemp -d)"; then
  echo "mktemp failed to make a temporary directory; exiting"
  exit 1
fi
image_file_xml="${tmp_dir}/${image_file_name}.xml"
echo "Creating temporary XML fsimage file at $image_file_xml"
hdfs oiv -p XML -i "$image_file" -o "$image_file_xml"
echo "Created temporary XML fsimage file"

if ! hdfs dfs -mkdir -p "$remote_dir"; then
  echo "Unable to create ${remote_dir}; exiting"
  rm -rf "${tmp_dir}"
  exit 1
fi
for file in ${image_file} ${image_file_xml} ${image_file_md5} ${version_file}; do
  echo "Uploading ${file}..."
  if ! hdfs dfs -copyFromLocal -f "$file" "${remote_dir}"; then
    echo "Error while uploading ${file}; exiting"
      rm -rf "${tmp_dir}"
    exit 1
  fi
done
rm -rf "${tmp_dir}"
echo "Complete!"
