#!/usr/bin/env bash
set -euo pipefail

# Calculate cooldown for each retry
get_cooldown_duration() {
  local completed_retry_count=$1
  local initial_cooldown=$2
  local multiplier=$3
  local max_cooldown=$4

  if ! [[ "$completed_retry_count" =~ ^[0-9]+$ ]] || \
     ! [[ "$initial_cooldown" =~ ^[0-9]+$ ]] || \
     ! [[ "$multiplier" =~ ^[0-9]+(\.[0-9]+)?$ ]] || \
     ! [[ "$max_cooldown" =~ ^(-1|[0-9]+)$ ]]; then
    echo "Error: Invalid input to get_cooldown_duration ($completed_retry_count, $initial_cooldown, $multiplier, $max_cooldown)" >&2
    return 1
  fi

  local calculated_cooldown
  calculated_cooldown=$(echo "$initial_cooldown + ($completed_retry_count * $multiplier)" | bc)

  # Round down
  local rounded_cooldown=${calculated_cooldown%.*}

  if (( max_cooldown > 0 && rounded_cooldown > max_cooldown )); then
    echo "$max_cooldown"
  else
    echo "$rounded_cooldown"
  fi
}

# Extract a parameter from the YAML, return first numeric value or blank
extract_yaml_param() {
  local yaml_file=$1
  local key=$2
  yq e ".processors | to_entries | .[] | .value.parameters[] | select(.key == \"$key\") | .value" "$yaml_file" 2>/dev/null | grep -E '^[0-9]+(\.[0-9]+)?$' | head -n 1
}

# Compute total retry cooldown time based on YAML values
calculate_max_retry_time() {
  local yaml_file=$1

  local attempts
  local initial_cooldown
  local multiplier
  local max_cooldown

  attempts=$(extract_yaml_param "$yaml_file" "attempts")
  initial_cooldown=$(extract_yaml_param "$yaml_file" "cooldown_between_attempts_seconds")
  multiplier=$(extract_yaml_param "$yaml_file" "cooldown_multiplier")
  max_cooldown=$(extract_yaml_param "$yaml_file" "max_cooldown")

  # Defaults
  [[ "$attempts" =~ ^[0-9]+$ ]] || { echo 0; return; }
  [[ "$initial_cooldown" =~ ^[0-9]+$ ]] || initial_cooldown=0
  [[ "$multiplier" =~ ^[0-9]+(\.[0-9]+)?$ ]] || multiplier=1
  [[ "$max_cooldown" =~ ^[0-9]+$ ]] || max_cooldown=2147483

  local total_cooldown_duration=0

  for (( retry=1; retry < attempts; retry++ )); do
    local interval_duration
    interval_duration=$(get_cooldown_duration "$retry" "$initial_cooldown" "$multiplier" "$max_cooldown") || {
      echo "Error: Failed cooldown for retry #$retry in $yaml_file" >&2
      echo -1
      return 1
    }

    total_cooldown_duration=$((total_cooldown_duration + interval_duration))
  done

  echo "$total_cooldown_duration"
}

# Get the maximum TTL from all TTL values found in a YAML file
get_max_ttl() {
  local yaml_file=$1
  local current_max=$2

  local ttl_values
  ttl_values=$(yq e '.processors | to_entries | .[] | .value.parameters[] | select(.key == "ttl_seconds") | .value' "$yaml_file" 2>/dev/null | grep -E '^[0-9]+$')

  local ttl
  while IFS= read -r ttl; do
    if (( ttl > current_max )); then
      current_max=$ttl
    fi
  done <<< "$ttl_values"

  echo "$current_max"
}

### Main logic
lunar_timeout=${LUNAR_SPOE_PROCESSING_TIMEOUT_SEC:-}
max_overall_ttl=${LUNAR_SPOE_PROCESSING_TIMEOUT_SEC:-0}
max_overall_retry_ttl=${LUNAR_RETRY_REQUEST_TIMEOUT_SEC:-0}

find "${LUNAR_PROXY_FLOW_DIRECTORY:?Must set LUNAR_PROXY_FLOW_DIRECTORY}" -type f \( -name "*.yaml" -o -name "*.yml" \) | while IFS= read -r yaml_file; do
  file_max_ttl=$(get_max_ttl "$yaml_file" "$max_overall_ttl")
  if (( file_max_ttl >= max_overall_ttl )); then
    echo -n $((file_max_ttl + 20)) > "/run/s6/container_environment/LUNAR_SPOE_PROCESSING_TIMEOUT_SEC"
    max_overall_ttl=$file_max_ttl
  fi

  file_max_retry_ttl=$(calculate_max_retry_time "$yaml_file")
  if (( file_max_retry_ttl >= max_overall_retry_ttl )); then
    echo -n $(((file_max_retry_ttl + 20) * 1000)) > "/run/s6/container_environment/LUNAR_RETRY_REQUEST_TIMEOUT_SEC"
    max_overall_retry_ttl=$file_max_retry_ttl
  fi
done
