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

DEFAULT_PERIOD_US=100000

usage() {
  local script_name
  script_name="${0##*/}"
  cat <<EOF
Usage:
  $script_name show    --namespace <ns> --pod <pod-name> [--path <cgroup-path>]
  $script_name set     --namespace <ns> --pod <pod-name> [--cores <n>] [--cpu-weight <1-10000>] [--cpu-weight-nice <-20..19>] [--memory-high <bytes|max>] [--io-weight <value>] [--io-max <value>]
  $script_name restore --namespace <ns> --pod <pod-name> [--all] [--path <cgroup-path>]

Examples:
  ./$script_name show --namespace default --pod app-7f9c8d6b5c-abcde
  ./$script_name set --namespace default --pod app-7f9c8d6b5c-abcde --cores 8
  ./$script_name set --namespace default --pod app-7f9c8d6b5c-abcde --cpu-weight 50
  ./$script_name set --namespace default --pod app-7f9c8d6b5c-abcde --memory-high 8589934592
  ./$script_name set --namespace default --pod app-7f9c8d6b5c-abcde --io-max '8:0 rbps=10485760 wbps=10485760'
  ./$script_name restore --namespace default --pod app-7f9c8d6b5c-abcde
  ./$script_name restore --namespace default --pod app-7f9c8d6b5c-abcde --all

Notes:
  - show/set/restore all read or write cgroup files through ssh on the target host.
  - --pod must be the exact Pod name. If the Pod does not exist, kubectl reports the error.
  - The script resolves the Pod's node InternalIP and uses root@<ip> for ssh.
  - show prints the Pod's declared resources.limits.cpu and matching cpu.max before current cgroup state.
  - set changes only relatively safe temporary controls: cpu.max, cpu.weight, cpu.weight.nice, memory.high, io.weight, and io.max.
  - restore resets cpu.max from Pod CPU limit by default.
  - restore --all also resets tunables to cgroup v2 default/unlimited values: cpu.weight=100, cpu.weight.nice=0, memory.high=max, io.weight='default 100', and existing io.max device rows to max.
  - --namespace and --pod require local kubectl to resolve the Pod, node, UID, and resource limit.
EOF
}

die() {
  echo "error: $*" >&2
  exit 1
}

require_value() {
  local name="$1"
  local value="${2:-}"
  [[ -n "$value" ]] || die "$name requires a value"
}

shell_quote() {
  printf "%q" "$1"
}

uid_to_pod_id() {
  local uid="$1"
  printf "pod%s" "${uid//-/_}"
}

resolve_uid_with_kubectl() {
  local namespace="$1"
  local pod="$2"
  kubectl get pod -n "$namespace" "$pod" -o jsonpath='{.metadata.uid}'
}

resolve_node_name_with_kubectl() {
  local namespace="$1"
  local pod="$2"
  kubectl get pod -n "$namespace" "$pod" -o jsonpath='{.spec.nodeName}'
}

resolve_node_ip_with_kubectl() {
  local node="$1"
  local address_type="$2"
  kubectl get node "$node" -o jsonpath="{range .status.addresses[?(@.type==\"$address_type\")]}{.address}{\"\\n\"}{end}" |
    awk 'NF { print; exit }'
}

resolve_host_with_kubectl() {
  local namespace="$1"
  local pod="$2"
  local ssh_user="$3"
  local address_type="$4"
  local node address

  node="$(resolve_node_name_with_kubectl "$namespace" "$pod")"
  [[ -n "$node" ]] || die "Pod has no spec.nodeName: $namespace/$pod"

  address="$(resolve_node_ip_with_kubectl "$node" "$address_type")"
  if [[ -z "$address" ]]; then
    address="$node"
  fi

  if [[ -n "$ssh_user" ]]; then
    printf "%s@%s\n" "$ssh_user" "$address"
  else
    printf "%s\n" "$address"
  fi
}

read_pod_cpu_limit() {
  local namespace="$1"
  local pod="$2"
  local limit

  limit="$(kubectl get pod -n "$namespace" "$pod" -o jsonpath='{.spec.containers[?(@.name=="compute")].resources.limits.cpu}')"
  if [[ -z "$limit" ]]; then
    limit="$(kubectl get pod -n "$namespace" "$pod" -o jsonpath='{range .spec.containers[*]}{.resources.limits.cpu}{"\n"}{end}' | awk 'NF { print; exit }')"
  fi

  [[ -n "$limit" ]] || die "Pod has no resources.limits.cpu: $namespace/$pod"
  printf "%s\n" "$limit"
}

ssh_run() {
  local host="$1"
  local script="$2"
  ssh "$host" "bash -s" <<<"$script"
}

remote_path_prelude() {
  local path="$1"
  local pod_uid="$2"

  if [[ -n "$path" ]]; then
    local quoted_path
    quoted_path="$(shell_quote "$path")"
    cat <<EOF
set -euo pipefail
path=$quoted_path
[[ -d "\$path" ]] || { echo "cgroup path not found: \$path" >&2; exit 2; }
EOF
    return
  fi

  local pod_id
  pod_id="$(uid_to_pod_id "$pod_uid")"
  local quoted_pod_id
  quoted_pod_id="$(shell_quote "$pod_id")"
  cat <<EOF
set -euo pipefail
pod_id=$quoted_pod_id
path=\$(find /sys/fs/cgroup -type d -name "*\${pod_id}*" | head -n 1)
[[ -n "\$path" ]] || { echo "cgroup path not found for \$pod_id" >&2; exit 2; }
EOF
}

restore_cpu_max() {
  local host="$1"
  local path="$2"
  local pod_uid="$3"
  local cpu_max="$4"
  local quoted_cpu_max
  quoted_cpu_max="$(shell_quote "$cpu_max")"
  ssh_run "$host" "$(remote_path_prelude "$path" "$pod_uid")
cpu_max=$quoted_cpu_max
echo \"path: \$path\"
printf '%s\n' \"\$cpu_max\" > \"\$path/cpu.max\"
echo 'restored:'
echo 'cpu.max:'
cat \"\$path/cpu.max\""
}

restore_all_tunables() {
  local host="$1"
  local path="$2"
  local pod_uid="$3"
  local cpu_max="$4"
  local quoted_cpu_max
  quoted_cpu_max="$(shell_quote "$cpu_max")"
  ssh_run "$host" "$(remote_path_prelude "$path" "$pod_uid")
cpu_max=$quoted_cpu_max
write_if_exists() {
  local file=\"\$1\"
  local value=\"\$2\"
  [[ -e \"\$path/\$file\" ]] || return 0
  printf '%s\n' \"\$value\" > \"\$path/\$file\"
}
restore_io_max() {
  [[ -e \"\$path/io.max\" ]] || return 0
  local current
  current=\$(cat \"\$path/io.max\")
  [[ -n \"\$current\" ]] || return 0
  while read -r dev rest; do
    [[ -n \"\$dev\" ]] || continue
    printf '%s rbps=max wbps=max riops=max wiops=max\n' \"\$dev\" > \"\$path/io.max\"
  done <<<\"\$current\"
}
echo \"path: \$path\"
write_if_exists cpu.max \"\$cpu_max\"
write_if_exists cpu.weight 100
write_if_exists cpu.weight.nice 0
write_if_exists memory.high max
write_if_exists io.weight 'default 100'
restore_io_max
echo 'restored:'
for file in cpu.max cpu.weight cpu.weight.nice memory.high io.weight io.max; do
  [[ -e \"\$path/\$file\" ]] || continue
  echo \"\$file:\"
  cat \"\$path/\$file\"
done"
}

show_cpu() {
  local host="$1"
  local path="$2"
  local pod_uid="$3"
  local cpu_limit="$4"
  local original_value="$5"
  printf 'Kubernetes declared resources\n'
  printf '%-22s %-24s %s\n' 'FIELD' 'VALUE' 'DESCRIPTION'
  printf '%-22s %-24s %s\n' 'cpu limit' "$cpu_limit" 'Pod declared CPU limit from Kubernetes'
  printf '%-22s %-24s %s\n' 'cpu.max' "$original_value" 'cpu.max value converted from declared CPU limit'
  printf '\n'
  ssh_run "$host" "$(remote_path_prelude "$path" "$pod_uid")
print_section() {
  printf '\n%s\n' \"\$1\"
  printf '%-24s %-32s %s\n' 'FILE' 'VALUE' 'DESCRIPTION'
}
describe_file() {
  case \"\$1\" in
    cpu.max) echo 'CPU hard quota: max period, or max for unlimited' ;;
    cpu.weight) echo 'Relative CPU share under contention, 1..10000' ;;
    cpu.weight.nice) echo 'Nice-style CPU weight, -20..19' ;;
    cpu.stat) echo 'CPU usage and throttling counters' ;;
    cpu.pressure) echo 'CPU pressure stall information' ;;
    memory.current) echo 'Current memory usage in bytes' ;;
    memory.peak) echo 'Peak memory usage in bytes' ;;
    memory.high) echo 'Memory soft throttle threshold' ;;
    memory.max) echo 'Memory hard limit; low values can OOM' ;;
    memory.events) echo 'Memory high/max/OOM event counters' ;;
    memory.pressure) echo 'Memory pressure stall information' ;;
    io.weight) echo 'Relative block IO weight' ;;
    io.max) echo 'Block IO hard limits per device' ;;
    io.stat) echo 'Block IO usage counters per device' ;;
    io.pressure) echo 'IO pressure stall information' ;;
    pids.current) echo 'Current process/thread count' ;;
    pids.max) echo 'Process/thread count limit' ;;
    pids.events) echo 'PID limit event counters' ;;
    *) echo '' ;;
  esac
}
show_single() {
  local file=\"\$1\"
  if [[ -e \"\$path/\$file\" ]]; then
    local value
    value=\$(tr '\n' ' ' < \"\$path/\$file\" | sed 's/[[:space:]]\+$//')
    printf '%-24s %-32s %s\n' \"\$file\" \"\$value\" \"\$(describe_file \"\$file\")\"
  fi
}
show_block() {
  local file=\"\$1\"
  if [[ -e \"\$path/\$file\" ]]; then
    printf '%-24s %-32s %s\n' \"\$file\" '<multi-line>' \"\$(describe_file \"\$file\")\"
    sed 's/^/  /' \"\$path/\$file\"
  fi
}
echo \"path: \$path\"
print_section 'CPU'
show_single cpu.max
show_single cpu.weight
show_single cpu.weight.nice
show_block cpu.stat
show_block cpu.pressure
print_section 'Memory'
show_single memory.current
show_single memory.peak
show_single memory.high
show_single memory.max
show_block memory.events
show_block memory.pressure
print_section 'IO'
show_single io.weight
show_single io.max
show_block io.stat
show_block io.pressure
print_section 'PIDs'
show_single pids.current
show_single pids.max
show_block pids.events"
}

set_tunables() {
  local host="$1"
  local path="$2"
  local pod_uid="$3"
  local cpu_max="$4"
  local cpu_weight="$5"
  local cpu_weight_nice="$6"
  local memory_high="$7"
  local io_weight="$8"
  local io_max="$9"
  local quoted_cpu_max quoted_cpu_weight quoted_cpu_weight_nice quoted_memory_high quoted_io_weight quoted_io_max
  quoted_cpu_max="$(shell_quote "$cpu_max")"
  quoted_cpu_weight="$(shell_quote "$cpu_weight")"
  quoted_cpu_weight_nice="$(shell_quote "$cpu_weight_nice")"
  quoted_memory_high="$(shell_quote "$memory_high")"
  quoted_io_weight="$(shell_quote "$io_weight")"
  quoted_io_max="$(shell_quote "$io_max")"
  ssh_run "$host" "$(remote_path_prelude "$path" "$pod_uid")
cpu_max=$quoted_cpu_max
cpu_weight=$quoted_cpu_weight
cpu_weight_nice=$quoted_cpu_weight_nice
memory_high=$quoted_memory_high
io_weight=$quoted_io_weight
io_max=$quoted_io_max
write_if_set() {
  local file=\"\$1\"
  local value=\"\$2\"
  [[ -n \"\$value\" ]] || return 0
  [[ -e \"\$path/\$file\" ]] || { echo \"missing cgroup file: \$file\" >&2; exit 2; }
  printf '%s\n' \"\$value\" > \"\$path/\$file\"
}
show_if_set() {
  local file=\"\$1\"
  local value=\"\$2\"
  [[ -n \"\$value\" ]] || return 0
  echo \"\$file:\"
  cat \"\$path/\$file\"
}
echo \"path: \$path\"
write_if_set cpu.max \"\$cpu_max\"
write_if_set cpu.weight \"\$cpu_weight\"
write_if_set cpu.weight.nice \"\$cpu_weight_nice\"
write_if_set memory.high \"\$memory_high\"
write_if_set io.weight \"\$io_weight\"
write_if_set io.max \"\$io_max\"
show_if_set cpu.max \"\$cpu_max\"
show_if_set cpu.weight \"\$cpu_weight\"
show_if_set cpu.weight.nice \"\$cpu_weight_nice\"
show_if_set memory.high \"\$memory_high\"
show_if_set io.weight \"\$io_weight\"
show_if_set io.max \"\$io_max\""
}

quota_for_cores() {
  local cores="$1"
  awk -v cores="$cores" -v period="$DEFAULT_PERIOD_US" '
    BEGIN {
      if (cores !~ /^[0-9]+([.][0-9]+)?$/ || cores <= 0) {
        exit 1
      }
      printf "%.0f %d\n", cores * period, period
    }
  ' || die "--cores must be a positive number"
}

quota_for_k8s_cpu() {
  local cpu="$1"
  awk -v cpu="$cpu" -v period="$DEFAULT_PERIOD_US" '
    BEGIN {
      if (cpu == "0" || cpu == "0m") {
        exit 1
      }
      if (cpu ~ /^[0-9]+m$/) {
        cores = substr(cpu, 1, length(cpu) - 1) / 1000
      } else if (cpu ~ /^[0-9]+([.][0-9]+)?$/) {
        cores = cpu + 0
      } else {
        exit 1
      }
      if (cores <= 0) {
        exit 1
      }
      printf "%.0f %d\n", cores * period, period
    }
  ' || die "unsupported Kubernetes CPU quantity: $cpu"
}

action="${1:-}"
[[ -n "$action" ]] || { usage; exit 1; }
if [[ "$action" == "--help" || "$action" == "-h" ]]; then
  usage
  exit 0
fi
shift

host=""
path=""
namespace=""
pod=""
cores=""
cpu_weight=""
cpu_weight_nice=""
memory_high=""
io_weight=""
io_max=""
restore_all="false"
ssh_user="root"
node_address_type="InternalIP"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --ssh-user)
      require_value "$1" "${2:-}"
      ssh_user="$2"
      shift 2
      ;;
    --node-address-type)
      require_value "$1" "${2:-}"
      node_address_type="$2"
      shift 2
      ;;
    --path)
      require_value "$1" "${2:-}"
      path="$2"
      shift 2
      ;;
    --namespace|-n)
      require_value "$1" "${2:-}"
      namespace="$2"
      shift 2
      ;;
    --pod)
      require_value "$1" "${2:-}"
      pod="$2"
      shift 2
      ;;
    --cores)
      require_value "$1" "${2:-}"
      cores="$2"
      shift 2
      ;;
    --cpu-weight)
      require_value "$1" "${2:-}"
      cpu_weight="$2"
      shift 2
      ;;
    --cpu-weight-nice)
      require_value "$1" "${2:-}"
      cpu_weight_nice="$2"
      shift 2
      ;;
    --memory-high)
      require_value "$1" "${2:-}"
      memory_high="$2"
      shift 2
      ;;
    --io-weight)
      require_value "$1" "${2:-}"
      io_weight="$2"
      shift 2
      ;;
    --io-max)
      require_value "$1" "${2:-}"
      io_max="$2"
      shift 2
      ;;
    --all)
      restore_all="true"
      shift
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      die "unknown argument: $1"
      ;;
  esac
done

[[ "$action" == "show" || "$action" == "set" || "$action" == "restore" ]] || die "unknown action: $action"
[[ -n "$namespace" && -n "$pod" ]] || die "--namespace and --pod are required"
kubectl get pod -n "$namespace" "$pod" >/dev/null

host="$(resolve_host_with_kubectl "$namespace" "$pod" "$ssh_user" "$node_address_type")"

if [[ -z "$path" ]]; then
  pod_uid="$(resolve_uid_with_kubectl "$namespace" "$pod")"
else
  pod_uid=""
fi

if [[ "$action" == "set" ]]; then
  [[ "$restore_all" == "false" ]] || die "--all is only valid for restore"
  [[ -n "$cores$cpu_weight$cpu_weight_nice$memory_high$io_weight$io_max" ]] || die "set requires at least one tunable option"
elif [[ "$action" == "restore" ]]; then
  [[ -z "$cores$cpu_weight$cpu_weight_nice$memory_high$io_weight$io_max" ]] || die "--cores/--cpu-weight/--cpu-weight-nice/--memory-high/--io-weight/--io-max are only valid for set"
else
  [[ "$restore_all" == "false" ]] || die "--all is only valid for restore"
  [[ -z "$cores$cpu_weight$cpu_weight_nice$memory_high$io_weight$io_max" ]] || die "--cores/--cpu-weight/--cpu-weight-nice/--memory-high/--io-weight/--io-max are only valid for set"
fi

case "$action" in
  show)
    cpu_limit="$(read_pod_cpu_limit "$namespace" "$pod")"
    original_value="$(quota_for_k8s_cpu "$cpu_limit")"
    show_cpu "$host" "$path" "$pod_uid" "$cpu_limit" "$original_value"
    ;;
  set)
    cpu_max=""
    if [[ -n "$cores" ]]; then
      cpu_max="$(quota_for_cores "$cores")"
    fi
    set_tunables "$host" "$path" "$pod_uid" "$cpu_max" "$cpu_weight" "$cpu_weight_nice" "$memory_high" "$io_weight" "$io_max"
    ;;
  restore)
    cpu_limit="$(read_pod_cpu_limit "$namespace" "$pod")"
    original_value="$(quota_for_k8s_cpu "$cpu_limit")"
    echo "restored from pod cpu limit: $cpu_limit"
    if [[ "$restore_all" == "true" ]]; then
      restore_all_tunables "$host" "$path" "$pod_uid" "$original_value"
    else
      restore_cpu_max "$host" "$path" "$pod_uid" "$original_value"
    fi
    ;;
esac