#!/usr/bin/env bash

readonly DEBUG=${DEBUG:-false}
[[ "$DEBUG" == "true" || "$DEBUG" == "1" ]] && set -o verbose

set -o errexit
set -o nounset
set -o pipefail

[[ ${BASH_SOURCE[0]} != "${0}" ]] && (
  echo "do not source ${BASH_SOURCE[0]}" >&2
  exit 36
)

# ---------------------------
# Exit codes
# ---------------------------
readonly ERR_CONFIG_MISSING=10
readonly ERR_DOCKER_NOT_RUNNING=11
readonly ERR_INVALID_PULLSECRET=13
readonly ERR_INVALID_CONFIG=14
readonly ERR_WAIT_FOR_TIMEOUT=15

readonly EXIT_SUCCESS=0

# ---------------------------
# Defaults for config values
# ---------------------------
# Model
readonly DEFAULT_MODELS_REPOSITORY="quay.io/apheris/hub-apps"

readonly DEFAULT_MOCK_ENABLED="true"
readonly DEFAULT_MOCK_TAG="0.49.0-mock-by-file"
readonly DEFAULT_MOCK_PORT=""
readonly DEFAULT_MOCK_CAPABILITIES="inference,finetuning"

readonly DEFAULT_OF3_ENABLED="false"
readonly DEFAULT_OF3_TAG="0.49.0-openfold3-by-file"
readonly DEFAULT_OF3_PORT=""
readonly DEFAULT_OF3_CAPABILITIES="inference,finetuning"

readonly DEFAULT_B2_ENABLED="false"
readonly DEFAULT_B2_TAG="0.49.0-boltz2-by-file"
readonly DEFAULT_B2_PORT=""
readonly DEFAULT_B2_CAPABILITIES="inference"

readonly DEFAULT_PROTENIX_ENABLED="false"
readonly DEFAULT_PROTENIX_TAG="0.49.0-protenix-by-file"
readonly DEFAULT_PROTENIX_PORT=""
readonly DEFAULT_PROTENIX_CAPABILITIES="inference"

readonly DEFAULT_GPU="all"
readonly DEFAULT_SHM_SIZE="8G"

# Hub
readonly DEFAULT_HUB_ENABLED="true"
readonly DEFAULT_HUB_REPOSITORY="quay.io/apheris/hub"
readonly DEFAULT_HUB_TAG="1.3.1"
readonly DEFAULT_HUB_PORT=""
readonly DEFAULT_HUB_MSA_ENABLED="false"

# Auth
readonly DEFAULT_AUTH_ENABLED="false"
readonly DEFAULT_AUTH_DOMAIN="https://apheris-ai-dev.eu.auth0.com/"
readonly DEFAULT_AUTH_AUDIENCE="https://hub.fold.apheris.com/api"
readonly DEFAULT_AUTH_ISSUER=""
readonly DEFAULT_AUTH_EXTRA_SCOPES=""
readonly DEFAULT_AUTH_CLIENT_ID="pndJvZpK2u1uAbZ6DHkNFVnhak4j9Xwi"
readonly DEFAULT_AUTH_PROVIDER_TYPE=""
readonly DEFAULT_AUTH_BROWSER_URL=""

# CA certificate path
readonly DEFAULT_HUB_CA_CERT=""
readonly DEFAULT_HUB_FINETUNING_HEARTBEAT_TIMEOUT="5m"

# Storage defaults (already effectively used, just making them explicit)
readonly DEFAULT_INPUT_DIR="\$HOME/apheris-hub/input"
readonly DEFAULT_OUTPUT_DIR="\$HOME/apheris-hub/output"

# DB defaults
readonly DEFAULT_DB_DEPLOY="true" # if you want local postgres by default
readonly DEFAULT_DB_PORT=""       # empty = random port
readonly DEFAULT_DB_DSN=""        # must be provided by config if external

# Registry secret
readonly DEFAULT_PULL_SECRET="" # must be provided by config

# ---------------------------
# Constants
# ---------------------------
readonly DEFAULT_CONFIG_FILE="./config.yaml"

readonly LABEL="apheris.hub=true"

readonly DEFAULT_NETWORK="apheris-hub"

readonly HUB_CONTAINER="apheris-hub"
readonly POSTGRES_CONTAINER="apheris-hub-postgres"
readonly POSTGRES_VOLUME="apheris-hub-postgres-data"

# Address to bind random ports to when config does not specify one.
# For local-only access, keep 127.0.0.1.
# If you need the service accessible from other hosts, change to 0.0.0.0.
readonly RANDOM_BIND_ADDR="127.0.0.1"

# Docker command (can be overridden by environment: DOCKER_CMD="sudo docker" ...)
readonly DOCKER_CMD="${DOCKER_CMD:-docker}"

# ---------------------------
# Minimal YAML parsing
# ---------------------------
declare -A YAML_VARS=()

parse_yaml() {
  local file="$1"
  local -a key_stack=()
  local line raw trimmed indent value rest key path
  local level i

  while IFS= read -r line || [[ -n "$line" ]]; do
    raw="$line"

    # Strip comments (simple, but works for our config)
    raw="${raw%%#*}"

    # Skip empty / whitespace-only lines
    [[ "$raw" =~ ^[[:space:]]*$ ]] && continue

    # Compute indentation (assumes 2 spaces per level)
    trimmed="${raw#"${raw%%[![:space:]]*}"}"
    indent="${raw%"$trimmed"}"
    level=$((${#indent} / 2))
    line="$trimmed"

    # Split key and value on first colon
    key="${line%%:*}"
    rest="${line#*:}"

    # Normalize key
    key="${key%%[[:space:]]*}"
    key="${key%$'\r'}"

    # Update key stack at this level
    key_stack[level]="$key"
    i=$((level + 1))
    while ((i < ${#key_stack[@]})); do
      unset "key_stack[i]"
      ((i++))
    done

    # Trim value
    rest="${rest#$' '}"
    rest="${rest%$'\r'}"

    # Trim trailing spaces
    while [[ "$rest" =~ [[:space:]]$ ]]; do
      rest="${rest%[[:space:]]}"
    done

    value="$rest"

    # No value => just a parent mapping (e.g. "hub:"), nothing to store
    [[ -z "$value" ]] && continue

    # Strip surrounding quotes if present
    if [[ "$value" == \"*\" && "$value" == *\" ]]; then
      value="${value:1:${#value}-2}"
    elif [[ "$value" == \'*\' && "$value" == *\' ]]; then
      value="${value:1:${#value}-2}"
    fi

    # Build full dotted path: e.g. hub.auth.enabled
    path="${key_stack[0]}"
    for ((i = 1; i <= level; i++)); do
      [[ -n "${key_stack[$i]-}" ]] || continue
      path+=".${key_stack[$i]}"
    done

    YAML_VARS["$path"]="$value"
  done <"$file"
}

yaml_get() {
  local key="$1"
  local default="${2-}"
  if [[ ${YAML_VARS[$key]+_} ]]; then
    echo "${YAML_VARS[$key]}"
  else
    echo "$default"
  fi
}

yaml_get_or_default_if_empty() {
  local key="$1"
  local default="${2-}"
  local value
  value="$(yaml_get "$key" "$default")"
  if [[ -z "$value" ]]; then
    echo "$default"
  else
    echo "$value"
  fi
}

yaml_get_list() {
  local file="$1"
  local target="$2"
  local default="${3-}"
  local -a key_stack=()
  local line raw trimmed indent rest key path value
  local level i
  local target_level=-1
  local result=""

  while IFS= read -r line || [[ -n "$line" ]]; do
    raw="$line"
    raw="${raw%%#*}"
    [[ "$raw" =~ ^[[:space:]]*$ ]] && continue

    trimmed="${raw#"${raw%%[![:space:]]*}"}"
    indent="${raw%"$trimmed"}"
    level=$((${#indent} / 2))
    line="$trimmed"

    if [[ "$line" == "- "* ]]; then
      if (( target_level >= 0 && level == target_level + 1 )); then
        value="${line#- }"
        value="${value%$'\r'}"

        while [[ "$value" =~ [[:space:]]$ ]]; do
          value="${value%[[:space:]]}"
        done

        if [[ "$value" != *:* ]]; then
          if [[ "$value" == \"*\" && "$value" == *\" ]]; then
            value="${value:1:${#value}-2}"
          elif [[ "$value" == \'*\' && "$value" == *\' ]]; then
            value="${value:1:${#value}-2}"
          fi

          if [[ -n "$result" ]]; then
            result+=$'\n'
          fi
          result+="$value"
        fi
      fi
      continue
    fi

    if (( target_level >= 0 && level <= target_level )); then
      break
    fi

    key="${line%%:*}"
    rest="${line#*:}"

    key="${key%%[[:space:]]*}"
    key="${key%$'\r'}"

    key_stack[level]="$key"
    i=$((level + 1))
    while ((i < ${#key_stack[@]})); do
      unset "key_stack[i]"
      ((i++))
    done

    path="${key_stack[0]}"
    for ((i = 1; i <= level; i++)); do
      [[ -n "${key_stack[$i]-}" ]] || continue
      path+=".${key_stack[$i]}"
    done

    rest="${rest#$' '}"
    rest="${rest%$'\r'}"

    if [[ "$path" == "$target" && -z "$rest" ]]; then
      target_level=$level
    fi
  done <"$file"

  if [[ -n "$result" ]]; then
    echo "$result"
  else
    echo "$default"
  fi
}

# ---------------------------
# Small helpers
# ---------------------------
die() {
  local msg="$1"
  local code="${2:-1}"
  echo "Error (code ${code}): ${msg}" >&2
  exit "$code"
}

silence() {
  # Usage: silence <cmd> [args...]
  # Runs the command, captures all output, and:
  #   - prints it only on error (or if DEBUG=true/1)
  #   - returns the original exit code
  local cmd="$1"
  shift
  local output exit_code
  set +o errexit
  output=$("$cmd" "$@" 2>&1)
  exit_code=$?
  set -o errexit

  if [[ $exit_code -ne 0 || "$DEBUG" == "true" || "$DEBUG" == "1" ]]; then
    echo "[command failed: $cmd $*] (exit code: $exit_code)" >&2
    echo "$output" >&2
  fi

  return $exit_code
}

wait_for() {
  local cmd=$1
  local retry_count=0
  local max_retries=40

  while [[ $retry_count -lt $max_retries ]] && ! (eval "$cmd" &>/dev/null); do
    sleep 0.5
    retry_count=$((retry_count + 1))
  done

  if [[ $retry_count -ge $max_retries ]]; then
    die "Timed out waiting for: $cmd" "$ERR_WAIT_FOR_TIMEOUT"
  fi
}

needs_cmd() {
  command -v "$1" >/dev/null 2>&1 || die "Required command not found: $1" "$2"
}

bool() {
  [[ "${1,,}" == "true" ]] && echo "true" || echo "false"
}

join_lines_with_commas() {
  local value="$1"
  local result=""
  local line

  while IFS= read -r line; do
    [[ -n "$line" ]] || continue
    if [[ -n "$result" ]]; then
      result+=","
    fi
    result+="$line"
  done <<<"$value"

  printf '%s' "$result"
}

normalize_list() {
  local value="$1"
  local result=""
  local item
  local -a items=()

  value="${value//$'\n'/,}"
  IFS=',' read -r -a items <<<"$value"

  for item in "${items[@]}"; do
    item="${item#"${item%%[![:space:]]*}"}"
    item="${item%"${item##*[![:space:]]}"}"
    [[ -n "$item" ]] || continue
    if [[ -n "$result" ]]; then
      result+=$'\n'
    fi
    result+="$item"
  done

  printf '%s' "$result"
}

expand_path() {
  local p="$1"

  # Trim leading/trailing whitespace to avoid YAML parsing quirks
  p="${p#"${p%%[![:space:]]*}"}"
  p="${p%"${p##*[![:space:]]}"}"

  case "$p" in
    \~) p="$HOME" ;;
    \~/*) p="$HOME/${p#~/}" ;;
    "\$HOME") p="$HOME" ;;
    "\$HOME/"*) p="$HOME/${p#\$HOME/}" ;;
    "\${HOME}") p="$HOME" ;;
    "\${HOME}/"*) p="$HOME/${p#\$\{HOME\}/}" ;;
  esac
  if [[ "$p" == "$HOME/~" ]]; then
    p="$HOME"
  elif [[ "$p" == "$HOME/~/"* ]]; then
    p="$HOME/${p#"$HOME/~/"}"
  fi

  printf '%s' "$p"
}

require_abs_path() {
  local p="$1"
  local label="$2"
  [[ -z "$p" || "$p" == /* ]] || die "$label must be an absolute path: $p" "$ERR_INVALID_CONFIG"
}

# ---------------------------
# Args / usage
# ---------------------------
CONFIG_FILE="$DEFAULT_CONFIG_FILE"
CMD="run"

usage() {
  cat <<EOF
Usage: ${0##*/} [command] [options]

Commands:
  run                      Start hub + models as per config.
  cleanup                  Stop/remove all labeled containers (hub + models) and network (if not used by other containers).
  cleanup-postgres         Remove the local Postgres container (if managed by this script). WARNING: This deletes persisted DB state.
  cleanup-storage          Remove storage folders as per config.
  diagnose                 Print current configuration and status of relevant docker resources.

Options:
  -c, --config PATH        Config file (default: $DEFAULT_CONFIG_FILE)
  -h, --help               Show this help
EOF
}

parse_args() {
  if [[ $# -gt 0 && "$1" != "-"* ]]; then
    CMD="$1"
    shift
  fi

  while [[ $# -gt 0 ]]; do
    case "$1" in
      -c | --config)
        [[ -n "${2-}" && "${2:0:1}" != "-" ]] || die "Option $1 requires a path" "$ERR_INVALID_CONFIG"
        CONFIG_FILE="$2"
        shift
        ;;
      -c=* | --config=*)
        CONFIG_FILE="${1#*=}"
        ;;
      -h | --help)
        usage
        exit $EXIT_SUCCESS
        ;;
      *)
        die "Unknown option: $1. Please use --help for usage information." "$ERR_INVALID_CONFIG"
        ;;
    esac
    shift
  done
}

validate_args() {
  [[ -f "$CONFIG_FILE" ]] || die "Config file not found: $CONFIG_FILE" "$ERR_CONFIG_MISSING"
  CONFIG_FILE="$(cd "$(dirname "$CONFIG_FILE")" && pwd)/$(basename "$CONFIG_FILE")"
  [[ -f "$CONFIG_FILE" ]] || die "Config file not found after path resolution: $CONFIG_FILE" "$ERR_CONFIG_MISSING"

  needs_cmd "$DOCKER_CMD" "$ERR_DOCKER_NOT_RUNNING"

  # Check Docker daemon reachability and permissions
  local output exit_code
  set +o errexit
  output=$("$DOCKER_CMD" info 2>&1)
  exit_code=$?
  set -o errexit

  if [[ $exit_code -ne 0 ]]; then
    echo "Error: Failed to talk to Docker daemon using: $DOCKER_CMD" >&2
    echo "---- docker info output ----" >&2
    echo "$output" >&2
    echo "----------------------------" >&2

    if grep -qi "permission denied" <<<"$output"; then
      cat >&2 <<EOF
It looks like Docker is running, but this user/command does not have permission
to access the Docker daemon socket.

Possible fixes:
  - Ensure the user running this script is in the 'docker' group (and re-login), OR
  - Run this script with a different Docker command, for example:
        DOCKER_CMD="sudo docker" $0 "$@"

Current user: $USER
Current DOCKER_CMD: $DOCKER_CMD
EOF
    else
      cat >&2 <<EOF
Docker daemon is not running or not accessible.

Please:
  - Make sure the Docker service is running (e.g. 'systemctl status docker'), AND
  - Ensure '$DOCKER_CMD info' works from your shell before running this script.
EOF
    fi

    exit "$ERR_DOCKER_NOT_RUNNING"
  fi
}

# ---------------------------
# Config loading
# ---------------------------
load_config() {
  parse_yaml "$CONFIG_FILE"

  AUTH_ENABLED="$(bool "$(yaml_get 'hub.auth.enabled' "$DEFAULT_AUTH_ENABLED")")"
  AUTH_DOMAIN="$(yaml_get_or_default_if_empty 'hub.auth.domain' "$DEFAULT_AUTH_DOMAIN")"
  AUTH_AUDIENCE="$(yaml_get_or_default_if_empty 'hub.auth.audience' "$DEFAULT_AUTH_AUDIENCE")"
  AUTH_ISSUER="$(yaml_get_or_default_if_empty 'hub.auth.issuer' "$DEFAULT_AUTH_ISSUER")"
  AUTH_EXTRA_SCOPES="$(yaml_get_or_default_if_empty 'hub.auth.extraScopes' "$DEFAULT_AUTH_EXTRA_SCOPES")"
  AUTH_CLIENT_ID="$(yaml_get_or_default_if_empty 'hub.auth.clientId' "$DEFAULT_AUTH_CLIENT_ID")"
  AUTH_PROVIDER_TYPE="$(yaml_get_or_default_if_empty 'hub.auth.providerType' "$DEFAULT_AUTH_PROVIDER_TYPE")"
  AUTH_BROWSER_URL="$(yaml_get_or_default_if_empty 'hub.auth.browserUrl' "$DEFAULT_AUTH_BROWSER_URL")"

  HUB_CA_CERT="$(yaml_get 'hub.caCert' "$DEFAULT_HUB_CA_CERT")"
  [[ -n "$HUB_CA_CERT" ]] && HUB_CA_CERT="$(expand_path "$HUB_CA_CERT")"
  require_abs_path "$HUB_CA_CERT" "hub.caCert"

  INPUT_DIR="$(expand_path "$(yaml_get 'storage.input' "$DEFAULT_INPUT_DIR")")"
  OUTPUT_DIR="$(expand_path "$(yaml_get 'storage.output' "$DEFAULT_OUTPUT_DIR")")"

  NETWORK_NAME="$(yaml_get 'network' "$DEFAULT_NETWORK")"

  HUB_ENABLED="$(bool "$(yaml_get 'hub.enabled' "$DEFAULT_HUB_ENABLED")")"
  HUB_REPOSITORY="$(yaml_get 'hub.repository' "$DEFAULT_HUB_REPOSITORY")"
  HUB_TAG="$(yaml_get 'hub.tag' "$DEFAULT_HUB_TAG")"
  HUB_PORT="$(yaml_get 'hub.port' "$DEFAULT_HUB_PORT")"
  HUB_MSA_ENABLED="$(bool "$(yaml_get 'hub.msa.enabled' "$DEFAULT_HUB_MSA_ENABLED")")"
  HUB_FINETUNING_HEARTBEAT_TIMEOUT="$(yaml_get_or_default_if_empty 'hub.finetuningHeartbeatTimeout' "$DEFAULT_HUB_FINETUNING_HEARTBEAT_TIMEOUT")"

  DB_DEPLOY="$(bool "$(yaml_get 'db.deploy' "$DEFAULT_DB_DEPLOY")")"
  DB_DSN="$(yaml_get 'db.dsn' "$DEFAULT_DB_DSN")"
  DB_PORT="$(yaml_get 'db.port' "$DEFAULT_DB_PORT")"

  PULL_SECRET="$(yaml_get 'pullSecret' "$DEFAULT_PULL_SECRET")"

  MOCK_ENABLED="$(bool "$(yaml_get 'models.mock.enabled' "$DEFAULT_MOCK_ENABLED")")"
  MOCK_REPOSITORY="$(yaml_get 'models.mock.repository' "$DEFAULT_MODELS_REPOSITORY")"
  MOCK_TAG="$(yaml_get 'models.mock.tag' "$DEFAULT_MOCK_TAG")"
  MOCK_PORT="$(yaml_get 'models.mock.port' "$DEFAULT_MOCK_PORT")"
  MOCK_CAPABILITIES="$(normalize_list "$(yaml_get_list "$CONFIG_FILE" 'models.mock.capabilities' "$DEFAULT_MOCK_CAPABILITIES")")"
  MOCK_WEIGHTS_DIR="$(yaml_get 'models.mock.weightsDir' "")"
  [[ -n "$MOCK_WEIGHTS_DIR" ]] && MOCK_WEIGHTS_DIR="$(expand_path "$MOCK_WEIGHTS_DIR")"
  MOCK_WEIGHTS_CONFIG="$(yaml_get 'models.mock.weightsConfigFile' "")"
  [[ -n "$MOCK_WEIGHTS_CONFIG" ]] && MOCK_WEIGHTS_CONFIG="$(expand_path "$MOCK_WEIGHTS_CONFIG")"
  MOCK_WEIGHTS_ENV="$(yaml_get 'models.mock.weightsEnv' "")"

  OF3_ENABLED="$(bool "$(yaml_get 'models.openfold3.enabled' "$DEFAULT_OF3_ENABLED")")"
  OF3_REPOSITORY="$(yaml_get 'models.openfold3.repository' "$DEFAULT_MODELS_REPOSITORY")"
  OF3_TAG="$(yaml_get 'models.openfold3.tag' "$DEFAULT_OF3_TAG")"
  OF3_GPUS="$(yaml_get 'models.openfold3.gpus' "$DEFAULT_GPU")"
  OF3_SHM="$(yaml_get 'models.openfold3.shm_size' "$DEFAULT_SHM_SIZE")"
  OF3_PORT="$(yaml_get 'models.openfold3.port' "$DEFAULT_OF3_PORT")"
  OF3_CAPABILITIES="$(normalize_list "$(yaml_get_list "$CONFIG_FILE" 'models.openfold3.capabilities' "$DEFAULT_OF3_CAPABILITIES")")"
  OF3_WEIGHTS_DIR="$(yaml_get 'models.openfold3.weightsDir' "")"
  [[ -n "$OF3_WEIGHTS_DIR" ]] && OF3_WEIGHTS_DIR="$(expand_path "$OF3_WEIGHTS_DIR")"
  OF3_WEIGHTS_CONFIG="$(yaml_get 'models.openfold3.weightsConfigFile' "")"
  [[ -n "$OF3_WEIGHTS_CONFIG" ]] && OF3_WEIGHTS_CONFIG="$(expand_path "$OF3_WEIGHTS_CONFIG")"
  OF3_WEIGHTS_ENV="$(yaml_get 'models.openfold3.weightsEnv' "")"

  B2_ENABLED="$(bool "$(yaml_get 'models.boltz2.enabled' "$DEFAULT_B2_ENABLED")")"
  B2_REPOSITORY="$(yaml_get 'models.boltz2.repository' "$DEFAULT_MODELS_REPOSITORY")"
  B2_TAG="$(yaml_get 'models.boltz2.tag' "$DEFAULT_B2_TAG")"
  B2_GPUS="$(yaml_get 'models.boltz2.gpus' "$DEFAULT_GPU")"
  B2_SHM="$(yaml_get 'models.boltz2.shm_size' "$DEFAULT_SHM_SIZE")"
  B2_PORT="$(yaml_get 'models.boltz2.port' "$DEFAULT_B2_PORT")"
  B2_CAPABILITIES="$(normalize_list "$(yaml_get_list "$CONFIG_FILE" 'models.boltz2.capabilities' "$DEFAULT_B2_CAPABILITIES")")"
  B2_WEIGHTS_DIR="$(yaml_get 'models.boltz2.weightsDir' "")"
  [[ -n "$B2_WEIGHTS_DIR" ]] && B2_WEIGHTS_DIR="$(expand_path "$B2_WEIGHTS_DIR")"
  B2_WEIGHTS_CONFIG="$(yaml_get 'models.boltz2.weightsConfigFile' "")"
  [[ -n "$B2_WEIGHTS_CONFIG" ]] && B2_WEIGHTS_CONFIG="$(expand_path "$B2_WEIGHTS_CONFIG")"
  B2_WEIGHTS_ENV="$(yaml_get 'models.boltz2.weightsEnv' "")"

  PROTENIX_ENABLED="$(bool "$(yaml_get 'models.protenix.enabled' "$DEFAULT_PROTENIX_ENABLED")")"
  PROTENIX_REPOSITORY="$(yaml_get 'models.protenix.repository' "$DEFAULT_MODELS_REPOSITORY")"
  PROTENIX_TAG="$(yaml_get 'models.protenix.tag' "$DEFAULT_PROTENIX_TAG")"
  PROTENIX_GPUS="$(yaml_get 'models.protenix.gpus' "$DEFAULT_GPU")"
  PROTENIX_SHM="$(yaml_get 'models.protenix.shm_size' "$DEFAULT_SHM_SIZE")"
  PROTENIX_PORT="$(yaml_get 'models.protenix.port' "$DEFAULT_PROTENIX_PORT")"
  PROTENIX_CAPABILITIES="$(normalize_list "$(yaml_get_list "$CONFIG_FILE" 'models.protenix.capabilities' "$DEFAULT_PROTENIX_CAPABILITIES")")"

  validate_config
}

validate_model() {
  local key="$1"
  local enabled="$2"
  local repository="$3"
  local tag="$4"
  local port="${5-}"

  [[ "$enabled" == "true" ]] || return 0

  [[ -n "$repository" ]] || die "models.${key}.repository must be set" "$ERR_INVALID_CONFIG"
  [[ -n "$tag" ]] || die "models.${key}.tag must be set" "$ERR_INVALID_CONFIG"

  if [[ -n "$port" ]]; then
    [[ "$port" =~ ^[0-9]+$ ]] || die "configured port for '${key}' is invalid: '$port'" "$ERR_INVALID_CONFIG"
  fi
}

validate_hub() {
  local enabled="$1"
  local repository="$2"
  local tag="$3"
  local port="${4-}"

  [[ "$enabled" == "true" ]] || return 0

  [[ -n "$repository" ]] || die "hub.repository must be set" "$ERR_INVALID_CONFIG"
  [[ -n "$tag" ]] || die "hub.tag must be set" "$ERR_INVALID_CONFIG"

  if [[ -n "$port" ]]; then
    [[ "$port" =~ ^[0-9]+$ ]] || die "configured port for 'hub' is invalid: '$port'" "$ERR_INVALID_CONFIG"
  fi
}

validate_db() {
  local deploy="$1"
  local dsn="$2"
  local port="${3-}"

  if [[ "$deploy" == "false" && -z "$dsn" ]]; then
    die "db.deploy=false requires db.dsn" "$ERR_INVALID_CONFIG"
  fi

  if [[ "$deploy" == "true" ]]; then
    if [[ -n "$port" ]]; then
      [[ "$port" =~ ^[0-9]+$ ]] || die "configured port for 'db' is invalid: '$port'" "$ERR_INVALID_CONFIG"
    fi
  fi
}

validate_model_weights() {
  local key="$1"
  local enabled="$2"
  local weights_dir="$3"
  local weights_config="$4"
  local weights_env="$5"

  [[ "$enabled" == "true" ]] || return 0

  # If no custom weights config, nothing to validate
  [[ -z "$weights_dir" && -z "$weights_config" && -z "$weights_env" ]] && return 0

  # With inline configuration, validate weightsDir when provided because it may still be mounted.
  if [[ -n "$weights_env" ]]; then
    if [[ -n "$weights_dir" ]]; then
      if [[ ! -d "$weights_dir" ]]; then
        die "models.${key}.weightsDir does not exist: $weights_dir" "$ERR_INVALID_CONFIG"
      fi
      require_abs_path "$weights_dir" "models.${key}.weightsDir"
    fi
    return 0
  fi

  # If using file-based configuration, validate both weightsDir and weightsConfigFile
  if [[ -n "$weights_dir" || -n "$weights_config" ]]; then
    if [[ -z "$weights_dir" ]]; then
      die "models.${key}.weightsDir must be set when using weightsConfigFile" "$ERR_INVALID_CONFIG"
    fi
    if [[ -z "$weights_config" ]]; then
      die "models.${key}.weightsConfigFile must be set when using weightsDir" "$ERR_INVALID_CONFIG"
    fi

    # Validate paths exist
    if [[ ! -d "$weights_dir" ]]; then
      die "models.${key}.weightsDir does not exist: $weights_dir" "$ERR_INVALID_CONFIG"
    fi
    if [[ ! -f "$weights_config" ]]; then
      die "models.${key}.weightsConfigFile does not exist: $weights_config" "$ERR_INVALID_CONFIG"
    fi

    # Validate paths are absolute
    require_abs_path "$weights_dir" "models.${key}.weightsDir"
    require_abs_path "$weights_config" "models.${key}.weightsConfigFile"
  fi
}

validate_capabilities() {
  local key="$1"
  local capabilities="$2"
  local supported_capabilities="$3"
  local cap

  while IFS= read -r cap; do
    [[ -n "$cap" ]] || continue
    if ! grep -Fxq "$cap" <<<"$supported_capabilities"; then
      die "models.${key}.capabilities contains unsupported capability: ${cap}. Supported values: ${supported_capabilities//$'\n'/, }" "$ERR_INVALID_CONFIG"
    fi
  done <<<"$capabilities"
}

validate_config() {
  validate_hub "$HUB_ENABLED" "$HUB_REPOSITORY" "$HUB_TAG" "${HUB_PORT-}"
  validate_db "$DB_DEPLOY" "$DB_DSN" "${DB_PORT-}"

  validate_model "mock" "$MOCK_ENABLED" "$MOCK_REPOSITORY" "$MOCK_TAG" "${MOCK_PORT-}"
  validate_model "openfold3" "$OF3_ENABLED" "$OF3_REPOSITORY" "$OF3_TAG" "${OF3_PORT-}"
  validate_model "boltz2" "$B2_ENABLED" "$B2_REPOSITORY" "$B2_TAG" "${B2_PORT-}"
  validate_model "protenix" "$PROTENIX_ENABLED" "$PROTENIX_REPOSITORY" "$PROTENIX_TAG" "${PROTENIX_PORT-}"

  validate_model_weights "mock" "$MOCK_ENABLED" "$MOCK_WEIGHTS_DIR" "$MOCK_WEIGHTS_CONFIG" "$MOCK_WEIGHTS_ENV"
  validate_model_weights "openfold3" "$OF3_ENABLED" "$OF3_WEIGHTS_DIR" "$OF3_WEIGHTS_CONFIG" "$OF3_WEIGHTS_ENV"
  validate_model_weights "boltz2" "$B2_ENABLED" "$B2_WEIGHTS_DIR" "$B2_WEIGHTS_CONFIG" "$B2_WEIGHTS_ENV"

  validate_capabilities "mock" "$MOCK_CAPABILITIES" "$(normalize_list "inference,finetuning")"
  validate_capabilities "openfold3" "$OF3_CAPABILITIES" "$(normalize_list "inference,finetuning")"
  validate_capabilities "boltz2" "$B2_CAPABILITIES" "$(normalize_list "inference")"
  validate_capabilities "protenix" "$PROTENIX_CAPABILITIES" "$(normalize_list "inference")"
}

# ---------------------------
# Core actions
# ---------------------------

# Helper: add custom weights configuration to docker run options
# Parameters: weights_env, weights_dir, weights_config
add_custom_weights_if_present() {
  local weights_env="$1"
  local weights_dir="$2"
  local weights_config="$3"

  CUSTOM_WEIGHTS_OPTS=()

  # weightsEnv takes precedence over file-based configuration
  if [[ -n "$weights_env" ]]; then
    CUSTOM_WEIGHTS_OPTS+=(-e "APH_AVAILABLE_WEIGHTS=${weights_env}")
    if [[ -n "$weights_dir" ]]; then
      CUSTOM_WEIGHTS_OPTS+=(--volume "${weights_dir}:/weights:ro")
    fi
  elif [[ -n "$weights_dir" && -n "$weights_config" ]]; then
    CUSTOM_WEIGHTS_OPTS+=(--volume "${weights_dir}:/weights:ro")
    CUSTOM_WEIGHTS_OPTS+=(--volume "${weights_config}:/config/weights.json:ro")
    CUSTOM_WEIGHTS_OPTS+=(-e "APH_WEIGHTS_CONFIG_FILE=/config/weights.json")
  fi
}

# Helper: run docker command, capture output, and handle errors
# Returns successfully if docker command succeeds, otherwise processes error and dies
run_docker_with_error_handling() {
  local name="$1"
  local component="$2"
  local image="$3"
  local port_display="$4" # e.g., "port 7777" or "with a random host port"
  shift 4
  # Remaining args are the full docker run command arguments

  local output
  # Use || true to prevent errexit from stopping execution on failure
  output=$("$DOCKER_CMD" run "$@" "$image" 2>&1) || {
    local exit_code=$?
    echo "[command failed: docker run ... $image] (exit code: $exit_code)" >&2
    echo "$output" >&2
    silence "$DOCKER_CMD" rm -f "$name" >/dev/null 2>&1 || true

    # Check for specific error types
    if echo "$output" | grep -qiE "(unauthorized|authentication required|access.*not authorized|pull access denied)"; then
      die "Failed to pull image '${image}': unauthorized. Please configure a valid 'pullSecret' in your config file." "$ERR_INVALID_PULLSECRET"
    elif echo "$output" | grep -qiE "(address already in use|bind.*failed)"; then
      die "Failed to start ${component} ${port_display} is already in use" "$ERR_INVALID_CONFIG"
    else
      die "Failed to start ${component} ${port_display}" "$ERR_INVALID_CONFIG"
    fi
  }
}

# Helper: start container with 2 modes:
# - If host_port is non-empty (configured in config), try to bind exactly that port once.
#   On failure, exit with clear error (no fallback).
# - If host_port is empty, bind RANDOM_BIND_ADDR:0:inner_port (random free port),
#   then discover and return the assigned host port.
start_container_with_port_choice() {
  local name="$1"
  local component="$2"
  local host_port="$3"
  local inner_port="$4"
  local image="$5"
  shift 5
  local extra_opts=("$@")

  # Ensure no stale container with this name
  silence "$DOCKER_CMD" rm -f "$name" >/dev/null 2>&1 || true

  if [[ -n "$host_port" ]]; then
    run_docker_with_error_handling "$name" "$component" "$image" "on configured port: port ${host_port}" \
      -d --name "$name" "${extra_opts[@]}" -p "${host_port}:${inner_port}"
  else
    run_docker_with_error_handling "$name" "$component" "$image" "with a random host port:" \
      -d --name "$name" "${extra_opts[@]}" -p "${RANDOM_BIND_ADDR}:0:${inner_port}"

    local mapped
    mapped="$("$DOCKER_CMD" port "$name" "$inner_port" 2>/dev/null | tail -n1 || true)"
    if [[ -z "$mapped" ]]; then
      die "Could not determine mapped host port for ${component}" "$ERR_INVALID_CONFIG"
    fi
    host_port="${mapped##*:}"
  fi

  echo "$host_port"
}

docker_login_if_needed() {
  if [[ -z "$PULL_SECRET" || "$PULL_SECRET" == "..." ]]; then
    echo "pullSecret not set — skipping docker login." >&2
    return 0
  fi

  local decoded
  if ! decoded="$(echo "$PULL_SECRET" | base64 --decode 2>/dev/null)"; then
    die "pullSecret is not valid base64" "$ERR_INVALID_PULLSECRET"
  fi

  local user pass
  user="${decoded%%:*}"
  pass="${decoded#*:}"

  if [[ -z "$user" || -z "$pass" || "$user" == "$pass" ]]; then
    die "pullSecret must decode to 'user:token'" "$ERR_INVALID_PULLSECRET"
  fi

  echo "Logging in to quay.io as robot user '${user}'..."
  echo -n "$pass" | silence "$DOCKER_CMD" login quay.io -u "$user" --password-stdin \
    || die "docker login failed" "$ERR_INVALID_PULLSECRET"
}

ensure_folders() {
  echo "Ensuring storage folders..."
  mkdir -p "$INPUT_DIR" "$OUTPUT_DIR"
  chmod 777 "$INPUT_DIR" "$OUTPUT_DIR" || true
}

ensure_network() {
  echo "Ensuring docker network '$NETWORK_NAME'..."

  if silence "$DOCKER_CMD" network inspect "$NETWORK_NAME" >/dev/null; then
    return 0
  fi

  echo "Network '$NETWORK_NAME' does not exist, trying to create it..." >&2
  if ! silence "$DOCKER_CMD" network create --driver bridge "$NETWORK_NAME" >/dev/null; then
    die "Failed to ensure docker network '$NETWORK_NAME'. Check Docker daemon and permissions." "$ERR_DOCKER_NOT_RUNNING"
  fi
}

stop_labeled_containers_except_postgres() {
  echo "Stopping labeled ApherisFold containers (models + hub)..."

  local ids
  ids="$("$DOCKER_CMD" ps -a --filter "label=$LABEL" --quiet || true)"

  if [[ -z "$ids" ]]; then
    echo "No labeled ApherisFold containers found."
    return 0
  fi

  while IFS= read -r id; do
    [[ -n "$id" ]] || continue
    local name
    name="$("$DOCKER_CMD" inspect --format '{{.Name}}' "$id" 2>/dev/null || echo "")"
    name="${name#/}"

    if [[ "$name" == "$POSTGRES_CONTAINER" ]]; then
      continue
    fi

    echo "Stopping container '$name' (ID: $id)..."
    silence "$DOCKER_CMD" stop "$id" >/dev/null || true
    echo "Removing container '$name' (ID: $id)..."
    silence "$DOCKER_CMD" rm "$id" >/dev/null || true
  done <<<"$ids"

  echo "Removed all labeled ApherisFold containers (except Postgres)."
}

cleanup_postgres() {
  # Explicit destructive action only.
  if [[ "$DB_DEPLOY" == "true" ]]; then
    echo "Removing local Postgres container '${POSTGRES_CONTAINER}' (this deletes persisted DB state)..."
    silence "$DOCKER_CMD" rm -f "$POSTGRES_CONTAINER" >/dev/null 2>&1 || true

    echo "Removing local Postgres data volume '$POSTGRES_VOLUME'..."
    silence "$DOCKER_CMD" volume rm "$POSTGRES_VOLUME" >/dev/null 2>&1 || true

    echo "Local Postgres and its data volume removed."
  else
    echo "No local Postgres to remove (db.deploy=false means using external database)."
  fi
}

cleanup_storage() {
  echo "Removing storage folders (this deletes all input/output data)..."

  if [[ -d "$INPUT_DIR" ]]; then
    echo "Removing input directory '$INPUT_DIR'..."
    rm -rf "$INPUT_DIR" || true
  else
    echo "Input directory does not exist: $INPUT_DIR"
  fi

  if [[ -d "$OUTPUT_DIR" ]]; then
    echo "Removing output directory '$OUTPUT_DIR'..."
    rm -rf "$OUTPUT_DIR" || true
  else
    echo "Output directory does not exist: $OUTPUT_DIR"
  fi

  echo "Storage folders removed."
}

start_postgres_if_needed() {
  if [[ "$DB_DEPLOY" != "true" ]]; then
    echo "Using external Postgres (dsn from config)."
    return 0
  fi

  local image="postgres:15-alpine"

  echo "Ensuring Postgres data volume '$POSTGRES_VOLUME'..."
  silence "$DOCKER_CMD" volume create "$POSTGRES_VOLUME" >/dev/null 2>&1 || true

  local opts=(
    --restart unless-stopped
    --network "$NETWORK_NAME"
    --volume "${POSTGRES_VOLUME}:/var/lib/postgresql/data"
    -e POSTGRES_USER=postgres
    -e POSTGRES_PASSWORD=postgres
    -e POSTGRES_DB=hub
  )

  echo "Starting local Postgres container '${POSTGRES_CONTAINER}'..."
  DB_PORT="$(start_container_with_port_choice "$POSTGRES_CONTAINER" "Postgres" "${DB_PORT-}" 5432 "$image" "${opts[@]}")"

  echo "Waiting for Postgres to be ready..."
  wait_for "$DOCKER_CMD exec $POSTGRES_CONTAINER pg_isready -U postgres"

  DB_DSN="postgres://postgres:postgres@${POSTGRES_CONTAINER}:5432/hub?sslmode=disable"
  echo "Local Postgres is ready on port ${DB_PORT}"
}

run_model_mock() {
  if [[ "$MOCK_ENABLED" != "true" ]]; then
    return 0
  fi

  local image="${MOCK_REPOSITORY}:${MOCK_TAG}"
  local opts=(
    --restart unless-stopped
    --network "$NETWORK_NAME"
    --volume "${INPUT_DIR}:/input"
    --volume "${OUTPUT_DIR}:/output"
    -e BASE_INPUT_DIRECTORY=/input
    -e BASE_OUTPUT_DIRECTORY=/output
    -e APH_MODEL_SCOPE="$(join_lines_with_commas "$MOCK_CAPABILITIES")"
  )

  add_custom_weights_if_present "$MOCK_WEIGHTS_ENV" "$MOCK_WEIGHTS_DIR" "$MOCK_WEIGHTS_CONFIG"
  opts+=("${CUSTOM_WEIGHTS_OPTS[@]}")

  echo "Starting mock model container..."
  MOCK_PORT="$(start_container_with_port_choice "mock" "mock model" "${MOCK_PORT-}" 8000 "$image" "${opts[@]}")"
  echo "Mock model container port: $MOCK_PORT"
}

run_model_openfold3() {
  if [[ "$OF3_ENABLED" != "true" ]]; then
    return 0
  fi

  local image="${OF3_REPOSITORY}:${OF3_TAG}"
  local opts=(
    --restart unless-stopped
    --gpus "$OF3_GPUS"
    --shm-size="$OF3_SHM"
    --network "$NETWORK_NAME"
    --volume "${INPUT_DIR}:/input"
    --volume "${OUTPUT_DIR}:/output"
    -e BASE_INPUT_DIRECTORY=/input
    -e BASE_OUTPUT_DIRECTORY=/output
    -e APH_MODEL_SCOPE="$(join_lines_with_commas "$OF3_CAPABILITIES")"
  )

  add_custom_weights_if_present "$OF3_WEIGHTS_ENV" "$OF3_WEIGHTS_DIR" "$OF3_WEIGHTS_CONFIG"
  opts+=("${CUSTOM_WEIGHTS_OPTS[@]}")

  echo "Starting openfold3 model container..."
  OF3_PORT="$(start_container_with_port_choice "openfold3" "openfold3 model" "${OF3_PORT-}" 8000 "$image" "${opts[@]}")"
  echo "Openfold3 model container port: $OF3_PORT"
}

run_model_boltz2() {
  if [[ "$B2_ENABLED" != "true" ]]; then
    return 0
  fi

  local image="${B2_REPOSITORY}:${B2_TAG}"
  local opts=(
    --restart unless-stopped
    --gpus "$B2_GPUS"
    --shm-size="$B2_SHM"
    --network "$NETWORK_NAME"
    --volume "${INPUT_DIR}:/input"
    --volume "${OUTPUT_DIR}:/output"
    -e BASE_INPUT_DIRECTORY=/input
    -e BASE_OUTPUT_DIRECTORY=/output
    -e APH_MODEL_SCOPE="$(join_lines_with_commas "$B2_CAPABILITIES")"
  )

  add_custom_weights_if_present "$B2_WEIGHTS_ENV" "$B2_WEIGHTS_DIR" "$B2_WEIGHTS_CONFIG"
  opts+=("${CUSTOM_WEIGHTS_OPTS[@]}")

  echo "Starting boltz2 model container..."
  B2_PORT="$(start_container_with_port_choice "boltz2" "boltz2 model" "${B2_PORT-}" 8000 "$image" "${opts[@]}")"
  echo "Boltz2 model container port: $B2_PORT"
}

run_model_protenix() {
  if [[ "$PROTENIX_ENABLED" != "true" ]]; then
    return 0
  fi

  local image="${PROTENIX_REPOSITORY}:${PROTENIX_TAG}"
  local opts=(
    --restart unless-stopped
    --gpus "$PROTENIX_GPUS"
    --shm-size="$PROTENIX_SHM"
    --network "$NETWORK_NAME"
    --volume "${INPUT_DIR}:/input"
    --volume "${OUTPUT_DIR}:/output"
    -e BASE_INPUT_DIRECTORY=/input
    -e BASE_OUTPUT_DIRECTORY=/output
    -e APH_MODEL_SCOPE="$(join_lines_with_commas "$PROTENIX_CAPABILITIES")"
  )

  echo "Starting protenix model container..."
  PROTENIX_PORT="$(start_container_with_port_choice "protenix" "protenix model" "${PROTENIX_PORT-}" 8000 "$image" "${opts[@]}")"
  echo "Protenix model container port: $PROTENIX_PORT"
}

run_hub() {
  if [[ "$HUB_ENABLED" != "true" ]]; then
    return 0
  fi

  local image="${HUB_REPOSITORY}:${HUB_TAG}"

  local opts=(
    --restart unless-stopped
    --network "$NETWORK_NAME"
    --volume "${INPUT_DIR}:/input"
    --volume "${OUTPUT_DIR}:/output"
    --volume "${CONFIG_FILE}:/apheris/hub-config.yaml:ro"
    -e APH_HUB_INPUT_ROOT_DIRECTORY=/input
    -e APH_HUB_OUTPUT_ROOT_DIRECTORY=/output
    -e APH_HUB_DATABASE_DSN="$DB_DSN"
    -e APH_HUB_CONFIG_FILE=/apheris/hub-config.yaml
    -e APH_HUB_AUTH_ENABLED="$AUTH_ENABLED"
    -e APH_HUB_AUTH_DOMAIN="$AUTH_DOMAIN"
    -e APH_HUB_AUTH_CLIENT_ID="$AUTH_CLIENT_ID"
    -e APH_HUB_AUTH_AUDIENCE="$AUTH_AUDIENCE"
    -e APH_HUB_AUTH_ISSUER="$AUTH_ISSUER"
    -e APH_HUB_AUTH_EXTRA_SCOPES="$AUTH_EXTRA_SCOPES"
    -e APH_HUB_MSA_ENABLED="$HUB_MSA_ENABLED"
  )

  if [[ -n "$AUTH_BROWSER_URL" ]]; then
    opts+=(-e APH_HUB_AUTH_BROWSER_URL="$AUTH_BROWSER_URL")
  fi
  if [[ -n "$AUTH_PROVIDER_TYPE" ]]; then
    opts+=(-e APH_HUB_AUTH_PROVIDER_TYPE="$AUTH_PROVIDER_TYPE")
  fi
  if [[ -n "$HUB_FINETUNING_HEARTBEAT_TIMEOUT" ]]; then
    opts+=(-e APH_HUB_FINETUNING_HEARTBEAT_TIMEOUT="$HUB_FINETUNING_HEARTBEAT_TIMEOUT")
  fi

  # Mount CA certificate if configured
  if [[ -n "$HUB_CA_CERT" ]]; then
    if [[ ! -f "$HUB_CA_CERT" ]]; then
      die "CA certificate file not found: $HUB_CA_CERT" "$ERR_INVALID_CONFIG"
    fi
    opts+=(--volume "${HUB_CA_CERT}:/etc/ssl/certs/custom-ca.crt:ro")
  fi

  # Pass through selected APH_HUB_* environment variables, for example to configure timeouts:
  #   APH_HUB_PENDING_REQUEST_TIMEOUT=6h ./deploy_apherisfold
  for var_name in \
    APH_HUB_PENDING_REQUEST_TIMEOUT \
    APH_HUB_ACCEPTED_REQUEST_TIMEOUT \
    APH_HUB_WEBSOCKET_ENABLE_CORS \
    APH_HUB_MSA_POLL_INTERVAL \
    APH_HUB_MSA_REQUEST_TIMEOUT \
    APH_HUB_MSA_SERVERS \
    APH_HUB_MSA_WORKER_POLL_INTERVAL \
    APH_HUB_MSA_WORKER_COUNT \
    APH_HUB_MSA_LEASE_DURATION \
    APH_HUB_MSA_MAX_ATTEMPTS \
    APH_HUB_MSA_RETRY_BASE_DELAY \
    APH_HUB_MSA_RETRY_MAX_DELAY; do
    local value="${!var_name-}"
    if [[ -n "$value" ]]; then
      opts+=(-e "${var_name}=${value}")
    fi
  done

  echo "Starting ApherisFold Hub container..."
  HUB_PORT="$(start_container_with_port_choice "$HUB_CONTAINER" "hub" "${HUB_PORT-}" 8080 "$image" "${opts[@]}")"
  echo "ApherisFold Hub container port: $HUB_PORT"
}

run_all() {
  docker_login_if_needed
  ensure_network
  ensure_folders

  stop_labeled_containers_except_postgres
  start_postgres_if_needed

  run_model_mock
  run_model_openfold3
  run_model_boltz2
  run_model_protenix
  run_hub

  echo ""
  echo "Running containers:"
  "$DOCKER_CMD" ps --filter "label=$LABEL"

  if [[ "$HUB_ENABLED" == "true" ]]; then
    if [[ -n "$HUB_PORT" ]]; then
      echo ""
      echo "ApherisFold Hub UI is running and accessible at:"
      echo "    http://localhost:${HUB_PORT}"
      echo ""
    else
      echo "Hub is enabled but no host port was determined."
      echo "Check 'docker ps' output above."
    fi
  fi
}

network_has_containers() {
  local net="$1"
  # If network doesn't exist, treat as "no containers"
  if ! silence "$DOCKER_CMD" network inspect "$net" >/dev/null 2>&1; then
    return 1
  fi

  local count
  count="$("$DOCKER_CMD" network inspect "$net" \
    --format '{{ len .Containers }}' 2>/dev/null || echo "0")"

  [[ "$count" -gt 0 ]]
}

cleanup_all() {
  stop_labeled_containers_except_postgres

  if network_has_containers "$NETWORK_NAME"; then
    echo "Keeping network '$NETWORK_NAME' because containers are still attached."
    echo "Attached containers:"
    "$DOCKER_CMD" network inspect "$NETWORK_NAME" --format '{{ range .Containers }}- {{ .Name }}{{"\n"}}{{ end }}' \
      2>/dev/null || true
  else
    echo "Removing docker network '$NETWORK_NAME' (no containers attached)..."
    silence "$DOCKER_CMD" network rm "$NETWORK_NAME" >/dev/null 2>&1 || true
  fi

  echo "Cleanup complete."
}

diagnose() {
  echo "== diagnose =="
  echo "config file:                        ${CONFIG_FILE}"
  echo ""

  echo "-- hub --"
  echo "hub.enabled:                        ${HUB_ENABLED}"
  echo "hub.repository:                     ${HUB_REPOSITORY}"
  echo "hub.tag:                            ${HUB_TAG}"
  echo "hub.port:                           $([[ -n "${HUB_PORT}" ]] && echo "${HUB_PORT}" || echo \(empty\))"
  echo "hub.finetuningHeartbeatTimeout:     $([[ -n "${HUB_FINETUNING_HEARTBEAT_TIMEOUT}" ]] && echo "${HUB_FINETUNING_HEARTBEAT_TIMEOUT}" || echo \(empty\))"
  echo "hub.caCert:                         $([[ -n "${HUB_CA_CERT}" ]] && echo "${HUB_CA_CERT}" || echo \(empty\))"
  echo ""
  echo "-- hub auth --"
  echo "hub.auth.enabled:                   ${AUTH_ENABLED}"
  echo "hub.auth.domain:                    ${AUTH_DOMAIN}"
  echo "hub.auth.audience:                  ${AUTH_AUDIENCE}"
  echo "hub.auth.issuer:                    $([[ -n "${AUTH_ISSUER}" ]] && echo "${AUTH_ISSUER}" || echo \(empty\))"
  echo "hub.auth.providerType:              $([[ -n "${AUTH_PROVIDER_TYPE}" ]] && echo "${AUTH_PROVIDER_TYPE}" || echo \(empty\))"
  echo "hub.auth.extraScopes:               $([[ -n "${AUTH_EXTRA_SCOPES}" ]] && echo "${AUTH_EXTRA_SCOPES}" || echo \(empty\))"
  echo "hub.auth.clientId:                  ${AUTH_CLIENT_ID}"
  echo "hub.auth.browserUrl:                $([[ -n "${AUTH_BROWSER_URL}" ]] && echo "${AUTH_BROWSER_URL}" || echo \(empty\))"
  echo ""
  echo "-- hub msa --"
  echo "hub.msa.enabled:                    ${HUB_MSA_ENABLED}"
  echo ""

  echo "-- storage --"
  echo "storage.input:                      ${INPUT_DIR}"
  echo "storage.output:                     ${OUTPUT_DIR}"
  echo ""

  echo "-- db --"
  echo "db.deploy:                          ${DB_DEPLOY}"
  echo "db.dsn:                             $([[ -n "${DB_DSN}" ]] && echo \(set\) || echo \(empty\))"
  echo "db.port:                            $([[ -n "${DB_PORT}" ]] && echo "${DB_PORT}" || echo \(empty\))"
  echo ""

  echo "-- models --"
  echo "models.mock.enabled:                ${MOCK_ENABLED}"
  echo "models.mock.repository:             ${MOCK_REPOSITORY}"
  echo "models.mock.tag:                    ${MOCK_TAG}"
  echo "models.mock.port:                   $([[ -n "${MOCK_PORT}" ]] && echo "${MOCK_PORT}" || echo \(empty\))"
  echo "models.mock.capabilities:           ${MOCK_CAPABILITIES//$'\n'/, }"
  echo "models.mock.weightsDir:             $([[ -n "${MOCK_WEIGHTS_DIR}" ]] && echo "${MOCK_WEIGHTS_DIR}" || echo \(empty\))"
  echo "models.mock.weightsConfigFile:      $([[ -n "${MOCK_WEIGHTS_CONFIG}" ]] && echo "${MOCK_WEIGHTS_CONFIG}" || echo \(empty\))"
  echo "models.mock.weightsEnv:             $([[ -n "${MOCK_WEIGHTS_ENV}" ]] && echo \(set\) || echo \(empty\))"
  echo ""
  echo "models.openfold3.enabled:           ${OF3_ENABLED}"
  echo "models.openfold3.repository:        ${OF3_REPOSITORY}"
  echo "models.openfold3.tag:               ${OF3_TAG}"
  echo "models.openfold3.port:              $([[ -n "${OF3_PORT}" ]] && echo "${OF3_PORT}" || echo \(empty\))"
  echo "models.openfold3.capabilities:      ${OF3_CAPABILITIES//$'\n'/, }"
  echo "models.openfold3.weightsDir:        $([[ -n "${OF3_WEIGHTS_DIR}" ]] && echo "${OF3_WEIGHTS_DIR}" || echo \(empty\))"
  echo "models.openfold3.weightsConfigFile: $([[ -n "${OF3_WEIGHTS_CONFIG}" ]] && echo "${OF3_WEIGHTS_CONFIG}" || echo \(empty\))"
  echo "models.openfold3.weightsEnv:        $([[ -n "${OF3_WEIGHTS_ENV}" ]] && echo \(set\) || echo \(empty\))"
  echo ""
  echo "models.boltz2.enabled:              ${B2_ENABLED}"
  echo "models.boltz2.repository:           ${B2_REPOSITORY}"
  echo "models.boltz2.tag:                  ${B2_TAG}"
  echo "models.boltz2.port:                 $([[ -n "${B2_PORT}" ]] && echo "${B2_PORT}" || echo \(empty\))"
  echo "models.boltz2.capabilities:         ${B2_CAPABILITIES//$'\n'/, }"
  echo "models.boltz2.weightsDir:           $([[ -n "${B2_WEIGHTS_DIR}" ]] && echo "${B2_WEIGHTS_DIR}" || echo \(empty\))"
  echo "models.boltz2.weightsConfigFile:    $([[ -n "${B2_WEIGHTS_CONFIG}" ]] && echo "${B2_WEIGHTS_CONFIG}" || echo \(empty\))"
  echo "models.boltz2.weightsEnv:           $([[ -n "${B2_WEIGHTS_ENV}" ]] && echo \(set\) || echo \(empty\))"
  echo ""
  echo "models.protenix.enabled:            ${PROTENIX_ENABLED}"
  echo "models.protenix.repository:         ${PROTENIX_REPOSITORY}"
  echo "models.protenix.tag:                ${PROTENIX_TAG}"
  echo "models.protenix.port:               $([[ -n "${PROTENIX_PORT}" ]] && echo "${PROTENIX_PORT}" || echo \(empty\))"
  echo "models.protenix.capabilities:       ${PROTENIX_CAPABILITIES//$'\n'/, }"
  echo ""

  echo "-- secrets --"
  echo "pullSecret (registry):              $([[ -z "${PULL_SECRET}" || "${PULL_SECRET}" == "..." ]] && echo \(empty\) || echo \(set\))"
  echo ""

  echo "-- network --"
  echo "network:                            ${NETWORK_NAME}"
  echo ""

  echo "-- docker version --"
  "$DOCKER_CMD" version || true
  echo ""

  echo "-- labeled containers --"
  "$DOCKER_CMD" ps -a --filter "label=${LABEL}" || true
  echo ""

  echo "-- network inspect --"
  "$DOCKER_CMD" network inspect "${NETWORK_NAME}" || true
  echo ""
}

# ---------------------------
# Main
# ---------------------------
parse_args "$@"
validate_args
load_config

case "$CMD" in
  run) run_all ;;
  cleanup) cleanup_all ;;
  cleanup-postgres) cleanup_postgres ;;
  cleanup-storage) cleanup_storage ;;
  diagnose) diagnose ;;
  *)
    die "Unknown command: $CMD. Please use --help for usage information." "$ERR_INVALID_CONFIG"
    ;;
esac