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

ROOT=$(CDPATH='' cd -- "$(dirname -- "$0")/.." && pwd)
MODE=${1:-compose}
OUTPUT=${2:-}
RPC_EXPRESSION='
  IO.puts(
    JSON.encode!(%{
      node: to_string(node()),
      port_count: :erlang.system_info(:port_count),
      port_limit: :erlang.system_info(:port_limit),
      process_count: :erlang.system_info(:process_count),
      memory_total: :erlang.memory(:total),
      memory_system: :erlang.memory(:system),
      port_table:
        :erlang.system_info(:allocated_areas)
        |> Keyword.fetch!(:port_table)
    })
  )
'

temporary=$(mktemp "${TMPDIR:-/tmp}/wnh-beam-runtime.XXXXXX")
trap 'rm -f "$temporary"' EXIT HUP INT TERM

record_probe() {
  local target=$1
  local component=$2
  local expected=$3
  local runtime_json=$4
  local cgroup_memory=$5
  local rss_kb=$6

  jq -cn \
    --arg target "$target" \
    --arg component "$component" \
    --argjson expected_port_limit "$expected" \
    --argjson cgroup_memory "$cgroup_memory" \
    --argjson rss_kb "$rss_kb" \
    --argjson runtime "$runtime_json" \
    '{
      target: $target,
      component: $component,
      expected_port_limit: $expected_port_limit,
      cgroup_memory: $cgroup_memory,
      rss_kb: $rss_kb
    } + $runtime' >>"$temporary"
}

case "$MODE" in
  compose)
    mapfile -t targets < <(
      docker compose -p "${COMPOSE_PROJECT_NAME:-who_need_help}" \
        ps -q web worker |
        sort -u
    )

    if [[ "${#targets[@]}" -eq 0 ]]; then
      echo "No running Compose web or worker replicas were found." >&2
      exit 1
    fi

    live_flags=$(
      for target in "${targets[@]}"; do
        docker inspect --format '{{range .Config.Env}}{{println .}}{{end}}' \
          "$target" |
          sed -n 's/^ERL_ZFLAGS=//p'
      done |
        sort -u
    )

    if [[ ! "$live_flags" =~ ^\+Q[[:space:]]+([0-9]+)$ ]]; then
      echo "The live Compose replicas do not expose one unambiguous +Q limit." >&2
      exit 1
    fi

    expected=${BASH_REMATCH[1]}

    for target in "${targets[@]}"; do
      component=$(
        docker inspect --format \
          '{{index .Config.Labels "com.docker.compose.service"}}' "$target"
      )
      runtime_json=$(
        docker exec "$target" /app/bin/who_need_help rpc "$RPC_EXPRESSION" |
          tail -n 1
      )
      cgroup_memory=$(docker exec "$target" cat /sys/fs/cgroup/memory.current)
      rss_kb=$(
        docker exec "$target" awk "/^VmRSS:/ {print \$2}" /proc/1/status
      )
      record_probe \
        "$target" "$component" "$expected" "$runtime_json" \
        "$cgroup_memory" "$rss_kb"
    done
    ;;

  kind)
    KUBECTL="$ROOT/.tools/bin/kubectl"
    CONTEXT=kind-who-need-help
    NAMESPACE=who-need-help

    if [[ ! -x "$KUBECTL" ]]; then
      echo "The project-owned kubectl is missing." >&2
      exit 1
    fi

    live_flags=$(
      "$KUBECTL" --context "$CONTEXT" --namespace "$NAMESPACE" \
        get deployment \
        who-need-help-who-need-help-web \
        who-need-help-who-need-help-worker \
        -o json |
        jq -r '
          [
            .items[].spec.template.spec.containers[].env[]
            | select(.name == "ERL_ZFLAGS")
            | .value
          ]
          | unique
          | if length == 1 then .[0] else empty end
        '
    )

    if [[ ! "$live_flags" =~ (^|[[:space:]])\+Q[[:space:]]+([0-9]+)($|[[:space:]]) ]]; then
      echo "The live deployments do not expose one unambiguous +Q limit." >&2
      exit 1
    fi

    expected=${BASH_REMATCH[2]}
    mapfile -t pods < <(
      "$KUBECTL" --context "$CONTEXT" --namespace "$NAMESPACE" \
        get pods \
        -l app.kubernetes.io/name=who-need-help \
        -o json |
        jq -r '
          .items[]
          | .metadata.labels["app.kubernetes.io/component"] as $component
          | select(
              ($component == "web" or $component == "worker") and
              .status.phase == "Running" and
              ([.status.containerStatuses[]?.ready] | all)
            )
          | [.metadata.name, $component]
          | @tsv
        '
    )

    desired=$(
      "$KUBECTL" --context "$CONTEXT" --namespace "$NAMESPACE" \
        get deployment \
        who-need-help-who-need-help-web \
        who-need-help-who-need-help-worker \
        -o json |
        jq '[.items[].spec.replicas] | add'
    )

    if [[ "${#pods[@]}" -ne "$desired" ]]; then
      echo "Expected $desired Ready application pods; observed ${#pods[@]}." >&2
      exit 1
    fi

    for entry in "${pods[@]}"; do
      IFS=$'\t' read -r pod component <<<"$entry"
      runtime_json=$(
        "$KUBECTL" --context "$CONTEXT" --namespace "$NAMESPACE" \
          exec "$pod" --container "$component" -- \
          /app/bin/who_need_help rpc "$RPC_EXPRESSION" |
          tail -n 1
      )
      cgroup_memory=$(
        "$KUBECTL" --context "$CONTEXT" --namespace "$NAMESPACE" \
          exec "$pod" --container "$component" -- \
          cat /sys/fs/cgroup/memory.current
      )
      rss_kb=$(
        "$KUBECTL" --context "$CONTEXT" --namespace "$NAMESPACE" \
          exec "$pod" --container "$component" -- \
          awk "/^VmRSS:/ {print \$2}" /proc/1/status
      )
      record_probe \
        "$pod" "$component" "$expected" "$runtime_json" \
        "$cgroup_memory" "$rss_kb"
    done
    ;;

  *)
    echo "Usage: $0 [compose|kind] [output.json]" >&2
    exit 1
    ;;
esac

summary=$(jq -s 'sort_by(.component, .target)' "$temporary")

if ! printf '%s\n' "$summary" |
  jq -e \
    --argjson expected "$expected" \
    'length > 0 and all(.port_limit == $expected)' >/dev/null; then
  echo "A running BEAM instance does not use the configured port limit." >&2
  exit 1
fi

if [[ -n "$OUTPUT" ]]; then
  printf '%s\n' "$summary" >"$OUTPUT"
fi

printf '%s\n' "$summary"