#!/bin/bash

set -euo pipefail

usage() {
  cat <<'EOF'
Usage: resolve-mirrored-image.sh --image <upstream-or-primary-image> [options]

Required:
  --image IMAGE                 Desired image reference to use locally.

Optional:
  --mirror-image IMAGE          Mirror image reference to check/pull/push.
  --upstream-image IMAGE        Upstream source image to import into mirror when mirror is missing.
  --registry-user USER          Registry username for mirror auth.
  --registry-password-env VAR   Environment variable name containing mirror registry password/token.
  --disable-remote-digest-check Skip remote digest freshness check for local images.
  --local-only                  Only verify local presence; do not pull.
  --quiet                       Suppress informational output except final image ref.

Behavior:
  1. If IMAGE exists locally and matches remote digest, use it.
  2. Otherwise try to pull IMAGE.
  3. If IMAGE pull fails and MIRROR_IMAGE is set, resolve via MIRROR_IMAGE.
  4. If MIRROR_IMAGE is missing and UPSTREAM_IMAGE is set, pull upstream, push mirror, then use mirror.
  5. Prints the resolved local image reference on stdout.
EOF
}

log() {
  if [[ "${QUIET}" != "true" ]]; then
    echo "$*" >&2
  fi
}

require_arg() {
  local value="$1"
  local name="$2"
  if [[ -z "${value}" ]]; then
    echo "Missing required argument: ${name}" >&2
    usage >&2
    exit 1
  fi
}

login_if_needed() {
  local image_ref="$1"
  local registry password_var password

  registry="${image_ref%%/*}"
  if [[ -z "${registry}" || "${registry}" == "${image_ref}" ]]; then
    return 0
  fi

  if [[ -z "${REGISTRY_USER}" || -z "${REGISTRY_PASSWORD_ENV}" ]]; then
    return 0
  fi

  password_var="${REGISTRY_PASSWORD_ENV}"
  password="${!password_var:-}"
  if [[ -z "${password}" ]]; then
    log "Skipping docker login for ${registry}: env ${password_var} is empty"
    return 0
  fi

  if [[ -n "${LOGGED_IN_REGISTRIES[${registry}]:-}" ]]; then
    return 0
  fi

  log "Logging into ${registry}"
  printf '%s' "${password}" | docker login "http://${registry}" -u "${REGISTRY_USER}" --password-stdin >/dev/null
  LOGGED_IN_REGISTRIES["${registry}"]=1
}

image_present_locally() {
  local image_ref="$1"
  docker image inspect "${image_ref}" >/dev/null 2>&1
}

is_digest_pinned_ref() {
  local image_ref="$1"
  [[ "${image_ref}" == *@sha256:* ]]
}

extract_remote_digest_buildx() {
  local image_ref="$1"
  local digest

  if ! digest="$(docker buildx imagetools inspect "${image_ref}" --format '{{json .Manifest.Digest}}' 2>/dev/null)"; then
    return 1
  fi

  digest="${digest//\"/}"
  digest="${digest//$'\n'/}"
  digest="${digest//$'\r'/}"
  digest="${digest// /}"

  if [[ "${digest}" == sha256:* ]]; then
    printf '%s\n' "${digest}"
    return 0
  fi

  return 1
}

extract_remote_digest_manifest() {
  local image_ref="$1"
  local digest

  if ! digest="$(docker manifest inspect "${image_ref}" 2>/dev/null | sed -n 's/.*"digest"[[:space:]]*:[[:space:]]*"\([^"]*\)".*/\1/p' | head -n 1)"; then
    return 1
  fi

  if [[ "${digest}" == sha256:* ]]; then
    printf '%s\n' "${digest}"
    return 0
  fi

  return 1
}

get_remote_digest() {
  local image_ref="$1"

  if is_digest_pinned_ref "${image_ref}"; then
    printf '%s\n' "${image_ref##*@}"
    return 0
  fi

  if extract_remote_digest_buildx "${image_ref}"; then
    return 0
  fi

  if extract_remote_digest_manifest "${image_ref}"; then
    return 0
  fi

  return 1
}

local_image_matches_digest() {
  local image_ref="$1"
  local digest="$2"

  docker image inspect --format '{{range .RepoDigests}}{{println .}}{{end}}' "${image_ref}" 2>/dev/null | grep -q "@${digest}$"
}

pull_image() {
  local image_ref="$1"
  login_if_needed "${image_ref}"
  docker pull "${image_ref}" >/dev/null
}

resolve_reference() {
  local image_ref="$1"
  local remote_digest

  if image_present_locally "${image_ref}"; then
    if [[ "${DISABLE_REMOTE_DIGEST_CHECK}" == "true" ]] || [[ "${LOCAL_ONLY}" == "true" ]] || is_digest_pinned_ref "${image_ref}"; then
      log "Using locally cached image ${image_ref}"
      return 0
    fi

    login_if_needed "${image_ref}"
    if remote_digest="$(get_remote_digest "${image_ref}")"; then
      if local_image_matches_digest "${image_ref}" "${remote_digest}"; then
        log "Using local image ${image_ref}; remote digest matches (${remote_digest})"
        return 0
      fi
      log "Local image ${image_ref} is stale; expected remote digest ${remote_digest}"
    else
      log "Unable to determine remote digest for ${image_ref}; attempting pull refresh"
    fi
  fi

  if [[ "${LOCAL_ONLY}" == "true" ]]; then
    return 1
  fi

  if pull_image "${image_ref}"; then
    log "Pulled image ${image_ref}"
    return 0
  fi

  return 1
}

mirror_available() {
  local image_ref="$1"
  if pull_image "${image_ref}"; then
    return 0
  fi
  return 1
}

promote_upstream_to_mirror() {
  require_arg "${UPSTREAM_IMAGE}" "--upstream-image"
  require_arg "${MIRROR_IMAGE}" "--mirror-image"

  log "Pulling upstream image ${UPSTREAM_IMAGE}"
  docker pull "${UPSTREAM_IMAGE}" >/dev/null

  login_if_needed "${MIRROR_IMAGE}"
  log "Tagging ${UPSTREAM_IMAGE} as ${MIRROR_IMAGE}"
  docker tag "${UPSTREAM_IMAGE}" "${MIRROR_IMAGE}"
  log "Pushing ${MIRROR_IMAGE}"
  docker push "${MIRROR_IMAGE}" >/dev/null
}

IMAGE=""
MIRROR_IMAGE=""
UPSTREAM_IMAGE=""
REGISTRY_USER=""
REGISTRY_PASSWORD_ENV=""
DISABLE_REMOTE_DIGEST_CHECK="false"
LOCAL_ONLY="false"
QUIET="false"
declare -A LOGGED_IN_REGISTRIES=()

while [[ $# -gt 0 ]]; do
  case "$1" in
    --image)
      IMAGE="${2:-}"
      shift 2
      ;;
    --mirror-image)
      MIRROR_IMAGE="${2:-}"
      shift 2
      ;;
    --upstream-image)
      UPSTREAM_IMAGE="${2:-}"
      shift 2
      ;;
    --registry-user)
      REGISTRY_USER="${2:-}"
      shift 2
      ;;
    --registry-password-env)
      REGISTRY_PASSWORD_ENV="${2:-}"
      shift 2
      ;;
    --disable-remote-digest-check)
      DISABLE_REMOTE_DIGEST_CHECK="true"
      shift
      ;;
    --local-only)
      LOCAL_ONLY="true"
      shift
      ;;
    --quiet)
      QUIET="true"
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown argument: $1" >&2
      usage >&2
      exit 1
      ;;
  esac
done

require_arg "${IMAGE}" "--image"

if resolve_reference "${IMAGE}"; then
  printf '%s\n' "${IMAGE}"
  exit 0
fi

if [[ "${LOCAL_ONLY}" == "true" ]]; then
  echo "Image unavailable via local-only mode: ${IMAGE}" >&2
  exit 1
fi

if [[ -n "${MIRROR_IMAGE}" ]]; then
  log "Primary image unavailable; checking mirror ${MIRROR_IMAGE}"
  if resolve_reference "${MIRROR_IMAGE}"; then
    printf '%s\n' "${MIRROR_IMAGE}"
    exit 0
  fi

  if [[ -n "${UPSTREAM_IMAGE}" ]]; then
    log "Mirror image ${MIRROR_IMAGE} missing; promoting from upstream ${UPSTREAM_IMAGE}"
    promote_upstream_to_mirror
    if ! pull_image "${MIRROR_IMAGE}"; then
      echo "Failed to pull promoted mirror image: ${MIRROR_IMAGE}" >&2
      exit 1
    fi
    printf '%s\n' "${MIRROR_IMAGE}"
    exit 0
  fi
fi

if [[ -n "${UPSTREAM_IMAGE}" ]]; then
  log "Falling back to upstream image ${UPSTREAM_IMAGE}"
  if resolve_reference "${UPSTREAM_IMAGE}"; then
    printf '%s\n' "${UPSTREAM_IMAGE}"
    exit 0
  fi
fi

echo "Failed to resolve image via primary/mirror/upstream path" >&2
echo "primary=${IMAGE}" >&2
if [[ -n "${MIRROR_IMAGE}" ]]; then
  echo "mirror=${MIRROR_IMAGE}" >&2
fi
if [[ -n "${UPSTREAM_IMAGE}" ]]; then
  echo "upstream=${UPSTREAM_IMAGE}" >&2
fi
exit 1