#!/bin/bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck disable=SC1091
. "$SCRIPT_DIR/common.sh"

usage() {
  cat >&2 <<'USAGE'
Usage:
  install_cse_worker.sh [--plan|--apply]

Distributes the cloud-engine tikv-worker and one rendered S3 config to every
TiKV host, installs cse-worker.service, and waits for /healthz. Defaults to plan.
USAGE
}

action="${1:---plan}"
case "$action" in
  --plan) apply=0 ;;
  --apply) apply=1 ;;
  --help) usage; exit 0 ;;
  *) usage; exit 1 ;;
esac

STATE_FILE="${STATE_FILE:-}"
require_env STATE_FILE
unset_vars TIKV_HOSTS TIUP_CLUSTER_DEPLOYED SSH_PORT CLUSTER_SSH_KEY CLUSTER_KNOWN_HOSTS \
  TIDB_RUNTIME_USER TIDB_RUNTIME_GROUP TIDB_RUNTIME_HOME WORKER_TARBALL WORKER_CONFIG
load_state_file "$STATE_FILE"

require_env \
  TIKV_HOSTS TIUP_CLUSTER_DEPLOYED SSH_PORT CLUSTER_SSH_KEY CLUSTER_KNOWN_HOSTS \
  TIDB_RUNTIME_USER TIDB_RUNTIME_GROUP TIDB_RUNTIME_HOME WORKER_TARBALL WORKER_CONFIG
[ "$TIUP_CLUSTER_DEPLOYED" = "1" ] || die "state does not record a deployed TiUP cluster"

[ -f "$CLUSTER_SSH_KEY" ] || die "missing cluster SSH key: $CLUSTER_SSH_KEY"
[ -f "$CLUSTER_KNOWN_HOSTS" ] || die "missing verified cluster known_hosts file: $CLUSTER_KNOWN_HOSTS"
require_port SSH_PORT
require_safe_name TIDB_RUNTIME_USER
require_safe_name TIDB_RUNTIME_GROUP
for name in CLUSTER_SSH_KEY CLUSTER_KNOWN_HOSTS TIDB_RUNTIME_HOME WORKER_TARBALL WORKER_CONFIG; do
  require_safe_absolute_path "$name"
done

if [ "$apply" -eq 0 ]; then
  cat <<EOF
PLAN ONLY: no remote worker files will be changed.
To apply on the control host, rerun: $0 --apply

Will install:
  worker tarball: $WORKER_TARBALL
  worker config:  $WORKER_CONFIG
  SSH user:       $TIDB_RUNTIME_USER
  SSH key:        $CLUSTER_SSH_KEY
  SSH port:       $SSH_PORT
  known hosts:    $CLUSTER_KNOWN_HOSTS
  runtime user:   $TIDB_RUNTIME_USER
  runtime group:  $TIDB_RUNTIME_GROUP
  runtime home:   $TIDB_RUNTIME_HOME
  TiKV hosts:     $TIKV_HOSTS
  state file:     $STATE_FILE
EOF
  exit 0
fi

need_cmd scp
need_cmd ssh
need_cmd tar
need_cmd mktemp
[ -f "$WORKER_TARBALL" ] || die "missing worker tarball: $WORKER_TARBALL"
[ -f "$WORKER_CONFIG" ] || die "missing worker config: $WORKER_CONFIG"

if grep -Eq 's3-(key-id|secret-key)' "$WORKER_CONFIG"; then
  die "worker config must use the instance profile and omit static S3 credentials"
fi

ssh_args=(-i "$CLUSTER_SSH_KEY" -p "$SSH_PORT" -o StrictHostKeyChecking=yes \
  -o "UserKnownHostsFile=$CLUSTER_KNOWN_HOSTS" -o ConnectTimeout=10 \
  -o ServerAliveInterval=15 -o ServerAliveCountMax=4)
scp_args=(-i "$CLUSTER_SSH_KEY" -P "$SSH_PORT" -o StrictHostKeyChecking=yes \
  -o "UserKnownHostsFile=$CLUSTER_KNOWN_HOSTS" -o ConnectTimeout=10)

tmp_dir="$(mktemp -d)"
trap 'rm -rf "$tmp_dir"' EXIT
tar -xzf "$WORKER_TARBALL" -C "$tmp_dir" tikv-worker
chmod +x "$tmp_dir/tikv-worker"

cat > "$tmp_dir/cse-worker.service" <<UNIT
[Unit]
Description=TiDB X Cloud Engine Worker
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
User=$TIDB_RUNTIME_USER
Group=$TIDB_RUNTIME_GROUP
ExecStart=$TIDB_RUNTIME_HOME/bin/tikv-worker --config $TIDB_RUNTIME_HOME/cse-worker.toml
Restart=on-failure
RestartSec=5
LimitNOFILE=1048576

[Install]
WantedBy=multi-user.target
UNIT

for host in $TIKV_HOSTS; do
  log "Installing cse-worker on $host"
  run_step "upload worker $host" 300 \
    scp "${scp_args[@]}" \
      "$tmp_dir/tikv-worker" "$WORKER_CONFIG" "$tmp_dir/cse-worker.service" \
      "$TIDB_RUNTIME_USER@$host:/tmp/"

  run_step "install worker $host" 300 \
    ssh "${ssh_args[@]}" "$TIDB_RUNTIME_USER@$host" \
    "set -e
     runtime_user=$TIDB_RUNTIME_USER
     runtime_group=$TIDB_RUNTIME_GROUP
     runtime_home=$TIDB_RUNTIME_HOME
     sudo systemctl stop cse-worker 2>/dev/null || true
     sudo install -d -o \"\$runtime_user\" -g \"\$runtime_group\" \"\$runtime_home/bin\" \"\$runtime_home/cse-worker-data\"
     sudo install -m 0755 -o \"\$runtime_user\" -g \"\$runtime_group\" /tmp/tikv-worker \"\$runtime_home/bin/tikv-worker\"
     sudo install -m 0644 -o \"\$runtime_user\" -g \"\$runtime_group\" /tmp/cse-worker.toml \"\$runtime_home/cse-worker.toml\"
     sudo install -m 0644 /tmp/cse-worker.service /etc/systemd/system/cse-worker.service
     sudo systemctl daemon-reload
     sudo systemctl enable --now cse-worker
     systemctl is-active cse-worker
     for _ in \$(seq 1 30); do
       curl -sf http://127.0.0.1:9500/healthz && exit 0
       sleep 1
     done
     sudo journalctl -u cse-worker --no-pager -n 80
     exit 1"
done

log "cse-worker installed and healthy on every TiKV host"