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

script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
# shellcheck source=scripts/cluster/common.sh
source "$script_dir/common.sh"

action=${1:-status}
[[ $action == status || $action == apply || $action == rollback ]] || {
  echo 'usage: configure-control-plane-io.sh {status|apply|rollback}' >&2
  exit 64
}

load_cluster_env
unit_name=easyai-k3s-io-scheduler.service
unit_path=/etc/systemd/system/$unit_name

cluster_ssh "$CLUSTER_HONGKONG_HOST" bash -s -- "$action" "$unit_name" "$unit_path" <<'REMOTE'
set -euo pipefail
action=$1
unit_name=$2
unit_path=$3

root_source=$(findmnt -n -o SOURCE /)
root_partition=${root_source#/dev/}
device=$(lsblk -no PKNAME "/dev/$root_partition" | head -n 1)
[[ $device =~ ^[a-zA-Z0-9._-]+$ && -f /sys/block/$device/queue/scheduler ]]

current_scheduler() {
  sed -n 's/.*\[\([^]]*\)\].*/\1/p' "/sys/block/$device/queue/scheduler"
}

case $action in
  status)
    configured=false
    systemctl is-enabled --quiet "$unit_name" 2>/dev/null && configured=true
    printf 'control_plane_io=STATUS node=hongkong device=%s scheduler=%s configured=%s\n' \
      "$device" "$(current_scheduler)" "$configured"
    ;;
  apply)
    grep -qw mq-deadline "/sys/block/$device/queue/scheduler" || {
      echo 'mq-deadline is not supported by the Hong Kong root block device' >&2
      exit 1
    }
    temporary=$(mktemp /etc/systemd/system/.easyai-k3s-io-scheduler.XXXXXX)
    cleanup() {
      [[ -f ${temporary:-} ]] && rm -f -- "$temporary" || true
    }
    trap cleanup EXIT
    cat >"$temporary" <<EOF
[Unit]
Description=EasyAI K3s control-plane disk scheduler
After=local-fs.target
Before=k3s.service

[Service]
Type=oneshot
ExecStart=/bin/sh -c 'echo mq-deadline > /sys/block/$device/queue/scheduler'
ExecStop=/bin/sh -c 'echo none > /sys/block/$device/queue/scheduler'
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
EOF
    chmod 0644 "$temporary"
    install -o root -g root -m 0644 "$temporary" "$unit_path"
    systemctl daemon-reload
    systemctl reset-failed "$unit_name" >/dev/null 2>&1 || true
    systemctl enable --now "$unit_name" >/dev/null
    [[ $(systemctl is-active "$unit_name") == active ]]
    [[ $(current_scheduler) == mq-deadline ]]
    printf 'control_plane_io=APPLIED node=hongkong device=%s scheduler=mq-deadline restart=false\n' "$device"
    ;;
  rollback)
    if [[ -f $unit_path ]]; then
      systemctl disable --now "$unit_name" >/dev/null 2>&1 || true
      rm -f -- "$unit_path"
      systemctl daemon-reload
    else
      printf '%s\n' none >"/sys/block/$device/queue/scheduler"
    fi
    [[ $(current_scheduler) == none ]]
    printf 'control_plane_io=ROLLED_BACK node=hongkong device=%s scheduler=none\n' "$device"
    ;;
esac
REMOTE