#!/bin/sh
set -eu

usage() {
  cat >&2 <<'EOF'
usage:
  remote-source-lock.sh acquire <lock-dir> <token> <lease-seconds> <wait-seconds> <owner>
  remote-source-lock.sh heartbeat <lock-dir> <token> <lease-seconds>
  remote-source-lock.sh check <lock-dir> <token>
  remote-source-lock.sh release <lock-dir> <token>
EOF
  exit 2
}

fail() {
  echo "remote source lock: $*" >&2
  exit 1
}

valid_token() {
  case "$1" in
    "" | *[!A-Za-z0-9._-]*) return 1 ;;
    *) return 0 ;;
  esac
}

valid_positive_integer() {
  case "$1" in
    "" | *[!0-9]* | 0) return 1 ;;
    *) return 0 ;;
  esac
}

now_epoch() {
  date +%s
}

read_value() (
  path=$1
  if [ -f "$path" ]; then
    sed -n '1p' "$path" 2>/dev/null || true
  fi
)

write_value() (
  path=$1
  value=$2
  temporary=$path.tmp.$$
  printf '%s\n' "$value" >"$temporary"
  mv "$temporary" "$path"
)

lease_is_expired() (
  lock_dir=$1
  now=$2
  expires=$(read_value "$lock_dir/lease-expires")
  case "$expires" in
    "" | *[!0-9]*) return 0 ;;
  esac
  [ "$expires" -le "$now" ]
)

metadata_is_malformed() (
  lock_dir=$1
  token=$(read_value "$lock_dir/token")
  expires=$(read_value "$lock_dir/lease-expires")
  valid_token "$token" || return 0
  valid_positive_integer "$expires" || return 0
  return 1
)

remove_if_malformed() (
  lock_dir=$1
  [ -d "$lock_dir" ] || return 0
  metadata_is_malformed "$lock_dir" || return 0
  tombstone=$lock_dir.malformed.$$
  mv "$lock_dir" "$tombstone" 2>/dev/null || return 0
  if metadata_is_malformed "$tombstone"; then
    rm -rf "$tombstone"
  elif [ ! -e "$lock_dir" ]; then
    mv "$tombstone" "$lock_dir" 2>/dev/null || true
  fi
)

remove_guard_if_owned() (
  guard=$1
  token=$2
  [ "$(read_value "$guard/token")" = "$token" ] || return 0
  tombstone=$guard.release.$token.$$
  if mv "$guard" "$tombstone" 2>/dev/null; then
    if [ "$(read_value "$tombstone/token")" = "$token" ]; then
      rm -rf "$tombstone"
    elif [ ! -e "$guard" ]; then
      mv "$tombstone" "$guard" 2>/dev/null || true
    fi
  fi
)

recover_expired_guard() (
  guard=$1
  now=$2
  [ -d "$guard" ] || return 0
  lease_is_expired "$guard" "$now" || return 0
  stale_token=$(read_value "$guard/token")
  valid_token "$stale_token" || return 0
  remove_guard_if_owned "$guard" "$stale_token"
)

acquire_guard() (
  guard=$1
  token=$2
  lease_seconds=$3
  wait_seconds=$4
  started=$(now_epoch)
  deadline=$((started + wait_seconds))
  malformed_since=
  while :; do
    now=$(now_epoch)
    if [ -d "$guard" ] && metadata_is_malformed "$guard"; then
      [ -n "$malformed_since" ] || malformed_since=$now
      if [ "$now" -ge "$((malformed_since + 2))" ]; then
        remove_if_malformed "$guard"
      fi
    else
      malformed_since=
      recover_expired_guard "$guard" "$now"
    fi
    if mkdir "$guard" 2>/dev/null; then
      write_value "$guard/token" "$token"
      write_value "$guard/lease-expires" "$((now + lease_seconds))"
      return 0
    fi
    if [ "$now" -ge "$deadline" ]; then
      fail "timed out waiting for lock guard $guard"
    fi
    sleep 1
  done
)

try_reclaim() (
  lock_dir=$1
  guard=$2
  guard_token=$3
  lease_seconds=$4
  now=$5

  if ! mkdir "$guard" 2>/dev/null; then
    recover_expired_guard "$guard" "$now"
    return 1
  fi
  write_value "$guard/token" "$guard_token"
  write_value "$guard/lease-expires" "$((now + lease_seconds))"

  stale_token=$(read_value "$lock_dir/token")
  if ! valid_token "$stale_token" || ! lease_is_expired "$lock_dir" "$now"; then
    remove_guard_if_owned "$guard" "$guard_token"
    return 1
  fi

  tombstone=$lock_dir.reclaimed.$stale_token.$$
  if ! mv "$lock_dir" "$tombstone" 2>/dev/null; then
    remove_guard_if_owned "$guard" "$guard_token"
    return 1
  fi
  if [ "$(read_value "$tombstone/token")" = "$stale_token" ] && lease_is_expired "$tombstone" "$now"; then
    rm -rf "$tombstone"
    remove_guard_if_owned "$guard" "$guard_token"
    return 0
  fi

  if [ ! -e "$lock_dir" ]; then
    mv "$tombstone" "$lock_dir" 2>/dev/null || true
  fi
  remove_guard_if_owned "$guard" "$guard_token"
  return 1
)

acquire() {
  [ "$#" -eq 5 ] || usage
  lock_dir=$1
  token=$2
  lease_seconds=$3
  wait_seconds=$4
  owner=$5
  valid_token "$token" || fail "ownership token must contain only letters, digits, dot, underscore, or dash"
  valid_positive_integer "$lease_seconds" || fail "lease seconds must be a positive integer"
  valid_positive_integer "$wait_seconds" || fail "wait seconds must be a positive integer"

  parent=$(dirname "$lock_dir")
  mkdir -p "$parent"
  guard=$lock_dir.reclaim
  started=$(now_epoch)
  deadline=$((started + wait_seconds))
  guard_token=$token.$$
  malformed_lock_since=
  malformed_guard_since=
  while :; do
    now=$(now_epoch)
    if [ -d "$guard" ] && metadata_is_malformed "$guard"; then
      [ -n "$malformed_guard_since" ] || malformed_guard_since=$now
      if [ "$now" -ge "$((malformed_guard_since + 2))" ]; then
        remove_if_malformed "$guard"
      fi
    else
      malformed_guard_since=
      recover_expired_guard "$guard" "$now"
    fi
    if [ ! -d "$guard" ] && mkdir "$lock_dir" 2>/dev/null; then
      write_value "$lock_dir/token" "$token"
      write_value "$lock_dir/lease-expires" "$((now + lease_seconds))"
      write_value "$lock_dir/owner" "$owner"
      printf '%s\n' "$token"
      return 0
    fi
    if [ -d "$lock_dir" ] && metadata_is_malformed "$lock_dir"; then
      [ -n "$malformed_lock_since" ] || malformed_lock_since=$now
      if [ "$now" -ge "$((malformed_lock_since + 2))" ] && [ ! -d "$guard" ] && mkdir "$guard" 2>/dev/null; then
        write_value "$guard/token" "$guard_token"
        write_value "$guard/lease-expires" "$((now + lease_seconds))"
        remove_if_malformed "$lock_dir"
        remove_guard_if_owned "$guard" "$guard_token"
      fi
    else
      malformed_lock_since=
    fi
    if [ -d "$lock_dir" ] && ! metadata_is_malformed "$lock_dir" && lease_is_expired "$lock_dir" "$now"; then
      try_reclaim "$lock_dir" "$guard" "$guard_token" "$lease_seconds" "$now" || true
    fi
    if [ "$now" -ge "$deadline" ]; then
      current_owner=$(read_value "$lock_dir/owner")
      current_expires=$(read_value "$lock_dir/lease-expires")
      fail "timed out waiting for live source lock $lock_dir (owner=${current_owner:-unknown}, lease_expires=${current_expires:-unknown})"
    fi
    sleep 1
  done
}

heartbeat() {
  [ "$#" -eq 3 ] || usage
  lock_dir=$1
  token=$2
  lease_seconds=$3
  valid_token "$token" || fail "invalid ownership token"
  valid_positive_integer "$lease_seconds" || fail "lease seconds must be a positive integer"
  guard=$lock_dir.reclaim
  guard_token=$token.heartbeat.$$
  acquire_guard "$guard" "$guard_token" 30 30
  if [ "$(read_value "$lock_dir/token")" != "$token" ]; then
    remove_guard_if_owned "$guard" "$guard_token"
    fail "ownership token mismatch for heartbeat"
  fi
  write_value "$lock_dir/lease-expires" "$(( $(now_epoch) + lease_seconds ))"
  remove_guard_if_owned "$guard" "$guard_token"
}

check() {
  [ "$#" -eq 2 ] || usage
  lock_dir=$1
  token=$2
  valid_token "$token" || fail "invalid ownership token"
  guard=$lock_dir.reclaim
  guard_token=$token.check.$$
  acquire_guard "$guard" "$guard_token" 30 30
  if [ "$(read_value "$lock_dir/token")" != "$token" ]; then
    remove_guard_if_owned "$guard" "$guard_token"
    fail "ownership token mismatch"
  fi
  if lease_is_expired "$lock_dir" "$(now_epoch)"; then
    remove_guard_if_owned "$guard" "$guard_token"
    fail "ownership lease expired"
  fi
  remove_guard_if_owned "$guard" "$guard_token"
}

release() {
  [ "$#" -eq 2 ] || usage
  lock_dir=$1
  token=$2
  valid_token "$token" || fail "invalid ownership token"
  guard=$lock_dir.reclaim
  guard_token=$token.release.$$
  acquire_guard "$guard" "$guard_token" 30 30
  if [ ! -d "$lock_dir" ]; then
    remove_guard_if_owned "$guard" "$guard_token"
    return 0
  fi
  if [ "$(read_value "$lock_dir/token")" != "$token" ]; then
    remove_guard_if_owned "$guard" "$guard_token"
    fail "ownership token mismatch for release"
  fi
  tombstone=$lock_dir.release.$token.$$
  if ! mv "$lock_dir" "$tombstone" 2>/dev/null; then
    remove_guard_if_owned "$guard" "$guard_token"
    fail "could not claim lock for release"
  fi
  if [ "$(read_value "$tombstone/token")" != "$token" ]; then
    if [ ! -e "$lock_dir" ]; then
      mv "$tombstone" "$lock_dir" 2>/dev/null || true
    fi
    remove_guard_if_owned "$guard" "$guard_token"
    fail "ownership token changed during release"
  fi
  rm -rf "$tombstone"
  remove_guard_if_owned "$guard" "$guard_token"
}

[ "$#" -ge 1 ] || usage
command=$1
shift
case "$command" in
  acquire) acquire "$@" ;;
  heartbeat) heartbeat "$@" ;;
  check) check "$@" ;;
  release) release "$@" ;;
  *) usage ;;
esac