#!/bin/bash
# Copyright UiPath 2021
#
# =================
# LICENSE AGREEMENT
# -----------------
#   Use of paid UiPath products and services is subject to the licensing agreement
#   executed between you and UiPath. Unless otherwise indicated by UiPath, use of free
#   UiPath products is subject to the associated licensing agreement available here:
#   https://www.uipath.com/legal/trust-and-security/legal-terms (or successor website).
#   You must not use this file separately from the product it is a part of or is associated with.

set -e
OPTIONAL_BUNDLE_PATH=${OPTIONAL_BUNDLE_PATH:-""}
TARGET_REGISTRY_URL=${TARGET_REGISTRY_URL:-""}
OFFLINE_BUNDLE_PATH=${OFFLINE_BUNDLE_PATH:-""}
TMP_OFFLINE_FOLDER=""
SF_FOLDER_NAME="sf-images"
EXTRACT_FOLDER=${EXTRACT_FOLDER:-""}
PREPEND_TAG="localhost:30071/"
MINIMUM_PODMAN_DISK_SIZE="150G"

declare -A image_map


function display_usage() {
  echo
  echo "***************************************************************************************"
  echo
  echo "hydrate-registry.sh hydrates a target OCI compliant registry from a downloaded images.tar bundle."
  echo
  echo "Usage:"
  echo "  $0 [flags]"
  echo
  echo "Parameters:"
  echo "  --target-registry-url                              Target OCI compliant registry for the artifacts."
  echo "  --offline-bundle-path                              Path to the offline bundle tar file."
  echo "  --optional-bundle-path                             Path to the optional bundle tar file."
  echo "  --extract-path                                     Location to extract the bundle."
  echo
  echo "Flags:"
  echo "  -h|--help                                         Display help"
  echo "  -d|--debug                                        Run in debug mode"
  echo
  echo "***************************************************************************************"
  echo
}

function parse_args() {
  while [ "$#" -gt 0 ]; do
    case "$1" in
      -d | --debug)
        echo "Running in debug"
        set -x
        shift
        ;;
      -h | --help)
        display_usage
        exit 0
        ;;
      --extract-path)
        if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
          EXTRACT_FOLDER="$2"
          shift 2
        else
          echo "Value for $1 is missing"
          exit 1
        fi
        ;;
      --offline-bundle-path)
        if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
          OFFLINE_BUNDLE_PATH="$2"
          shift 2
        else
          echo "Value for $1 is missing"
          exit 1
        fi
        ;;
      --optional-bundle-path)
        if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
          OPTIONAL_BUNDLE_PATH="$2"
          SF_FOLDER_NAME="sf-images"
          shift 2
        else
          echo "Value for $1 is missing"
          exit 1
        fi
        ;;
      --target-registry-url)
        if [ -n "$2" ] && [ "${2:0:1}" != "-" ]; then
          TARGET_REGISTRY_URL="$2"
          shift 2
        else
          echo "Value for $1 is missing"
          exit 1
        fi
        ;;
      *)
        echo "Invalid option: $1"
        exit 1
        ;;
    esac
  done

  if [ -z "$EXTRACT_FOLDER" ]; then
    echo "Mandatory option --extract-path is missing."
    display_usage
    exit 1
  fi

  # Check if mandatory options are provided
  if [ -z "$TARGET_REGISTRY_URL" ]; then
    echo "Mandatory option --target-registry-url is missing."
    display_usage
    exit 1
  fi

  if [ -n "$OFFLINE_BUNDLE_PATH" ] && [ -n "$OPTIONAL_BUNDLE_PATH" ]; then
    echo "Mandatory option --offline-bundle-path or --optional-bundle-path is missing. Only one should be provided."
    exit 1
  fi

  if [ -z "$OFFLINE_BUNDLE_PATH" ] && [ -z "$OPTIONAL_BUNDLE_PATH" ]; then
    echo "Mandatory option --offline-bundle-path or --optional-bundle-path is missing."
    exit 1
  fi
}

function validate_oci_compliance() {
  echo "Validating if the target registry is OCI compliant..."
  local chart_name="oci-compliance-checker"
  local temp_chart_dir="./${chart_name}"
  local timestamp
  local chart_version

  timestamp="$(date +%s)"
  chart_version="0.0.1-${timestamp}"

  helm create "${temp_chart_dir}" > /dev/null

  # Package the chart
  helm package "${temp_chart_dir}" --version "${chart_version}" -d . > /dev/null

  # Try to push the chart
  if helm push "${chart_name}-${chart_version}.tgz" "oci://${TARGET_REGISTRY_URL}/test" > /dev/null 2>&1; then
    echo "Registry is OCI compliant."
  else
    echo "Registry is not OCI compliant. Please check the target registry configuration."
    rm -rf "${temp_chart_dir}" "${chart_name}-${chart_version}.tgz"
    exit 1
  fi

  # Cleanup
  rm -rf "${temp_chart_dir}" "${chart_name}-${chart_version}.tgz"
}

function validate_podman_disk_size() {
  # Identify podman image storage location from 'podman info' command
  local podman_storage_path
  podman_storage_path=$(podman info --format '{{json .Store.GraphRoot}}') || {
    echo "[ERROR] podman info command failed. Please make sure that Podman is installed."
    exit 1
  }
  podman_storage_path=$(echo "${podman_storage_path}" | tr -d '"')

  local current_size_bytes
  local expected_size_bytes

  current_size_bytes=$(findmnt --target "${podman_storage_path}" --bytes -o SIZE -n)
  expected_size_bytes=$(numfmt --from=auto "${MINIMUM_PODMAN_DISK_SIZE}")

  if [[ "$current_size_bytes" -lt "$expected_size_bytes" ]]; then
    echo "Mount point '${podman_storage_path}' size($(numfmt --to=iec "$current_size_bytes") or ${current_size_bytes}) is less than expected($(numfmt --to=iec "$expected_size_bytes") or ${expected_size_bytes})"
    echo "[WARNING] ${MINIMUM_PODMAN_DISK_SIZE} of free disk space for Podman is a prerequisite for hydrating the registry."
    echo "You can change the storage path by updating 'graphRoot' in the podman configuration."

    while true; do
      printf "Do you wish to continue (y/n)? [y]: "
      read -r input
      case "${input}" in
      [Yy]*) break ;;
      [Nn]*)
        echo "Operation aborted."
        exit 1
        ;;
      "") break ;;
      *) printf "\nInvalid choice\n" ;;
      esac
    done
  fi
}

function extract_bundle() {
  if [ ! -e "${EXTRACT_FOLDER}" ]; then
    echo "Path does not exist: ${EXTRACT_FOLDER}"
    exit 1
  fi
  echo "Importing bundle at ${OFFLINE_BUNDLE_PATH}"
  TMP_OFFLINE_FOLDER="$(mktemp -d -p "${EXTRACT_FOLDER}" -t offline-bundle.XXXXXXXXXX)"
  echo "Extracting ${OFFLINE_BUNDLE_PATH} to ${TMP_OFFLINE_FOLDER}"
  tar -xzf "${OFFLINE_BUNDLE_PATH}" -C "${TMP_OFFLINE_FOLDER}"
  echo "Files used for the offline bundle have been extracted to ${TMP_OFFLINE_FOLDER}"
}

push_and_add_image() {
  local image_name="$1"
  echo "Tagging image ${PREPEND_TAG}${image_name} -> ${TARGET_REGISTRY_URL}/${image_name}"
  if ! podman tag "${PREPEND_TAG}${image_name}" "${TARGET_REGISTRY_URL}/${image_name}"; then
    echo "Failed to tag image ${image_name}"
    exit 1
  fi
  echo "Pushing image ${TARGET_REGISTRY_URL}/${image_name}"

  if ! podman push "${TARGET_REGISTRY_URL}/${image_name}" --tls-verify=false; then
    echo "Failed to push image ${image_name}"
    exit 1
  fi

  echo "Pushed image ${image_name} to ${TARGET_REGISTRY_URL}"
  image_map["${TARGET_REGISTRY_URL}/${image_name}"]="1"

  echo "Added image ${TARGET_REGISTRY_URL}/${image_name} to image set"
}

podman_load_and_push() {
  images_output=$(podman images --noheading)

  while read -r line; do
    # Skip the empty line, which is the first line
    if [[ "$line" = "" ]]; then
      continue
    fi

    # Extract the image names and tags, then add them to the map after checking if they exist
    image_name=$(echo "$line" | awk '{gsub(/^localhost:30071\//, ""); print $1}')
    image_tag=$(echo "$line" | awk '{print $2}')
    full_image_name="$image_name:$image_tag"
    if [ "$full_image_name" = "<none>:<none>" ]; then
        continue
    fi

    # Check if the image exists in the set before pushing
    if [[ -n "${image_map["$full_image_name"]}" ]]
    then
      echo "Image $full_image_name already loaded and pushed."
    else
      echo "Image $full_image_name loaded but is not pushed. Pushing..."

      # Push the image using podman push and add it to the set
      push_and_add_image "$full_image_name"
      image_map["$full_image_name"]="1"
    fi
  done <<< "$images_output"

  # Adding sleep to -
  # 1) to not run podman images too often as images are loaded after all its layers which takes a few seconds
  # 2) to not run podman push too often as it leads to 403 errors from registry
  sleep 5
}

function import_sf_images() {
  local dev_null="/dev/null"
  local image_inventory="${TMP_OFFLINE_FOLDER}/${SF_FOLDER_NAME}/image-tags.txt"
  local image_tar="${TMP_OFFLINE_FOLDER}/${SF_FOLDER_NAME}/images.tar"

  echo "Populating Docker Registry at https://${TARGET_REGISTRY_URL}"

  if [[ -f "${image_inventory}" ]]; then
    echo "Image inventory file $image_inventory found."
  else
    echo "Image inventory file $image_inventory not found."
    exit 1
  fi

  if [[ -f "${image_tar}" ]]; then
    echo "Image inventory file $image_tar found."
  else
    echo "Image inventory file $image_tar not found."
    exit 1
  fi

  local all_images_present=1

  if err=$(podman rmi -f -a 2>&1); then
    echo "Removed pre-loaded from podman, continuing..."
  else
    echo "Failed to remove pre-loaded images from podman: $err"
    exit 1
  fi

  while read -r image_name; do
    podman image exists "${TARGET_REGISTRY_URL}/${image_name}" || all_images_present=0
  done <"${image_inventory}"

  if ((all_images_present != 1)); then
    podman load -q -i "${image_tar}" &
    load_pid=$!

    while ps -p $load_pid > $dev_null; do
      podman_load_and_push
    done
  fi

  # Above loop will finish with podman load comannd. We need to check if some images were not pushed
  echo "Podman load finished, checking if some images were not pushed..."
  podman_load_and_push

  echo  "Populated Docker Registry"
}

function import_sf_charts() {
  local chart_inventory="${TMP_OFFLINE_FOLDER}/sf-helm/helm-tags.txt"

  echo "Populating Helm Registry at ${TARGET_REGISTRY_URL}"
  if [[ ! -f "${chart_inventory}" ]]; then
    echo "Chart inventory file ${chart_inventory} not found"
    exit 1
  fi
  while read -r line; do
    local chart_name file_name repo_url chart_namespace
    chart_name=$(awk '{print $1}' <<<"${line}")
    file_name=$(awk '{print $2}' <<<"${line}")
    local try=0
    local maxtry=10
    local status="notready"

    while [[ ${status} == "notready" ]] && ((try != maxtry)); do
      try=$((try + 1))
      chart_namespace=$(echo "${chart_name}" | cut -d "/" -f 1)
      repo_url="oci://${TARGET_REGISTRY_URL}/${chart_namespace}"
      echo "Pushing local chart ${TMP_OFFLINE_FOLDER}/sf-helm/${file_name} to ${repo_url}...try ${try}/${maxtry}"
      helm push "${TMP_OFFLINE_FOLDER}/sf-helm/${file_name}" "${repo_url}" && status="ready"
      [[ ${status} == "ready" ]] || sleep 2
    done

    if [[ ! ${status} == "ready" ]]; then
      echo "Failed to push chart ${chart_name} to ${repo_url}"
      exit 1
    fi
  done <"${chart_inventory}"

  echo "Populated Helm Registry"
}

function hydrate_from_bundle() {
  if [[ -z "${OPTIONAL_BUNDLE_PATH:-}" ]]; then
    import_sf_charts
  fi
  import_sf_images
}

function verify_podman_authentication() {
  if podman login "${TARGET_REGISTRY_URL}" --get-login &> /dev/null; then
    echo "Podman is logged in to ${TARGET_REGISTRY_URL}"
  else
    echo "Podman is not logged in. Please run 'podman login ${TARGET_REGISTRY_URL}' to log in."
    exit 1
  fi
}

function main() {
  validate_oci_compliance
  verify_podman_authentication
  validate_podman_disk_size
  if [[ -n "${OPTIONAL_BUNDLE_PATH:-}" ]]; then
    OFFLINE_BUNDLE_PATH="${OPTIONAL_BUNDLE_PATH}"
  fi
  extract_bundle
  hydrate_from_bundle

  echo "Cleaning up temporary files"
  rm -rf "${TMP_OFFLINE_FOLDER:?}/"*
}

parse_args "$@"
main