#!/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:-all}
[[ $action == backup || $action == install || $action == all ]] || {
  echo 'usage: bootstrap-k3s.sh [backup|install|all]' >&2
  exit 64
}

load_cluster_env
require_commands curl node openssl ssh scp shasum

secret_root="$cluster_root/.local-secrets/cluster"
cluster_environment="$secret_root/cluster.env"
kubeconfig="$secret_root/kubeconfig"
k3s_version=v1.36.2+k3s1
k3s_sha256=65a55ec56c24eab44383086166ec620a491952b7e23941a49ddca6e8a4c4b4de
k3s_binary="$secret_root/k3s-${k3s_version}-amd64"

install -d -m 0700 "$secret_root"

backup_legacy_k3s() {
  local timestamp remote_archive local_directory local_archive checksum object_prefix
  timestamp=$(date -u '+%Y%m%dT%H%M%SZ')
  remote_archive=/root/easyai-k3s-sqlite-pre-ha-"$timestamp".tar.gz
  local_directory=$(mktemp -d "$secret_root/backup.XXXXXX")
  local_archive=$local_directory/$(basename "$remote_archive")
  trap '[[ ${local_directory:-} == "$secret_root"/backup.* ]] && rm -rf -- "$local_directory"' RETURN

  echo '[k3s] archiving legacy SQLite state and staging manifests'
  cluster_ssh "$CLUSTER_NINGBO_HOST" bash -s -- "$remote_archive" <<'REMOTE'
set -euo pipefail
archive=$1
work_directory=$(mktemp -d /root/easyai-k3s-backup.XXXXXX)
trap '[[ $work_directory == /root/easyai-k3s-backup.* ]] && rm -rf -- "$work_directory"' EXIT
install -d -m 0700 "$work_directory/state"
if [[ -d /var/lib/rancher/k3s/server/db ]]; then
  cp -a /var/lib/rancher/k3s/server/db "$work_directory/state/"
fi
if [[ -d /var/lib/rancher/k3s/server/manifests ]]; then
  cp -a /var/lib/rancher/k3s/server/manifests "$work_directory/state/"
fi
if command -v k3s >/dev/null 2>&1 && systemctl is-active --quiet k3s; then
  k3s kubectl get namespace easyai-auth-center-staging -o yaml \
    >"$work_directory/easyai-auth-center-staging-namespace.yaml" 2>/dev/null || true
  k3s kubectl get all,configmap,secret,pvc,ingress \
    -n easyai-auth-center-staging -o yaml \
    >"$work_directory/easyai-auth-center-staging-resources.yaml" 2>/dev/null || true
  k3s kubectl get nodes -o wide >"$work_directory/nodes.txt"
  k3s --version >"$work_directory/k3s-version.txt"
fi
tar -C "$work_directory" -czf "$archive" .
chmod 0600 "$archive"
sha256sum "$archive" >"$archive.sha256"
REMOTE

  cluster_scp "$CLUSTER_NINGBO_HOST:$remote_archive" "$local_archive"
  cluster_scp "$CLUSTER_NINGBO_HOST:$remote_archive.sha256" "$local_archive.remote.sha256"
  checksum=$(shasum -a 256 "$local_archive" | awk '{print $1}')
  grep -q "^$checksum " "$local_archive.remote.sha256" || {
    echo 'legacy K3s archive checksum mismatch' >&2
    return 1
  }
  printf '%s  %s\n' "$checksum" "$(basename "$local_archive")" >"$local_archive.sha256"
  object_prefix="easyai-ai-gateway/production/pre-ha/${timestamp}"
  node "$script_dir/oss-object.mjs" put "$object_prefix/$(basename "$local_archive")" "$local_archive"
  node "$script_dir/oss-object.mjs" put "$object_prefix/$(basename "$local_archive").sha256" "$local_archive.sha256"
  node "$script_dir/oss-object.mjs" head "$object_prefix/$(basename "$local_archive")"
  cluster_ssh "$CLUSTER_NINGBO_HOST" \
    "rm -f -- '$remote_archive' '$remote_archive.sha256'"
  echo "legacy_k3s_backup=PASS object_prefix=$object_prefix sha256=$checksum"
}

write_node_config() {
  local output=$1
  local node_name=$2
  local node_ip=$3
  local mode=$4
  local taint=${5:-}
  umask 077
  {
    printf 'token-file: /etc/rancher/k3s/cluster-token\n'
    printf 'node-name: %s\n' "$node_name"
    printf 'node-ip: %s\n' "$node_ip"
    printf 'advertise-address: %s\n' "$node_ip"
    printf 'flannel-iface: wg0\n'
    printf 'write-kubeconfig-mode: "0600"\n'
    printf 'secrets-encryption: true\n'
    printf 'disable:\n  - traefik\n  - servicelb\n'
    printf 'tls-san:\n  - 10.77.0.1\n  - 10.77.0.2\n  - 10.77.0.3\n'
    printf 'etcd-arg:\n  - heartbeat-interval=500\n  - election-timeout=5000\n'
    printf 'etcd-snapshot-schedule-cron: "0 */6 * * *"\n'
    printf 'etcd-snapshot-retention: 28\n'
    printf 'etcd-snapshot-compress: true\n'
    printf 'etcd-s3: true\n'
    printf 'etcd-s3-config-secret: easyai-etcd-s3\n'
    if [[ $mode == init ]]; then
      printf 'cluster-init: true\n'
    else
      printf 'server: https://10.77.0.1:6443\n'
    fi
    if [[ -n $taint ]]; then
      printf 'node-taint:\n  - %s\n' "$taint"
    fi
  } >"$output"
  chmod 0600 "$output"
}

install_cluster() {
  local configuration_directory temporary_binary actual_k3s_sha256
  for host in "$CLUSTER_NINGBO_HOST" "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST"; do
    cluster_ssh "$host" 'ping -c 2 -W 2 10.77.0.1 >/dev/null'
  done

  if [[ ! -f $k3s_binary ]] ||
      [[ $(shasum -a 256 "$k3s_binary" | awk '{print $1}') != "$k3s_sha256" ]]; then
    temporary_binary=$(mktemp "$secret_root/.k3s-download.XXXXXX")
    trap 'rm -f -- "${temporary_binary:-}"' RETURN
    curl -fL --retry 3 --connect-timeout 15 \
      "https://github.com/k3s-io/k3s/releases/download/${k3s_version/+/%2B}/k3s" \
      -o "$temporary_binary"
    actual_k3s_sha256=$(shasum -a 256 "$temporary_binary" | awk '{print $1}')
    [[ $actual_k3s_sha256 == "$k3s_sha256" ]] || {
      echo "K3s binary checksum mismatch: $actual_k3s_sha256" >&2
      return 1
    }
    chmod 0700 "$temporary_binary"
    mv "$temporary_binary" "$k3s_binary"
    temporary_binary=
  fi
  actual_k3s_sha256=$(shasum -a 256 "$k3s_binary" | awk '{print $1}')
  [[ $actual_k3s_sha256 == "$k3s_sha256" ]]

  if [[ ! -f $cluster_environment ]]; then
    umask 077
    printf 'K3S_TOKEN=%s\n' "$(openssl rand -hex 32)" >"$cluster_environment"
    chmod 0600 "$cluster_environment"
  fi
  assert_private_file "$cluster_environment"
  # shellcheck source=/dev/null
  source "$cluster_environment"
  : "${K3S_TOKEN:?}"
  [[ $K3S_TOKEN =~ ^[0-9a-f]{64}$ ]] || {
    echo 'invalid local K3s token' >&2
    return 1
  }

  configuration_directory=$(mktemp -d "$secret_root/config.XXXXXX")
  trap '[[ ${configuration_directory:-} == "$secret_root"/config.* ]] && rm -rf -- "$configuration_directory"' RETURN
  write_node_config "$configuration_directory/ningbo.yaml" easyai-ningbo 10.77.0.1 init
  write_node_config "$configuration_directory/hongkong.yaml" easyai-hongkong 10.77.0.2 join
  write_node_config "$configuration_directory/los-angeles.yaml" easyai-los-angeles 10.77.0.3 join \
    'easyai.io/witness=true:NoSchedule'
  cp "$cluster_root/deploy/kubernetes/registries.yaml" \
    "$configuration_directory/registries.yaml"
  umask 077
  cat >"$configuration_directory/etcd-s3-secret.yaml" <<EOF
apiVersion: v1
kind: Secret
metadata:
  name: easyai-etcd-s3
  namespace: kube-system
type: etcd.k3s.cattle.io/s3-config-secret
stringData:
  etcd-s3-endpoint: ${ALI_REGION}.aliyuncs.com
  etcd-s3-access-key: ${ALI_KEY}
  etcd-s3-secret-key: ${ALI_SECRET}
  etcd-s3-bucket: ${ALI_BUCKET}
  etcd-s3-folder: easyai-ai-gateway/production/k3s-etcd
  etcd-s3-region: ${ALI_REGION}
  etcd-s3-bucket-lookup-type: dns
  etcd-s3-insecure: "false"
  etcd-s3-skip-ssl-verify: "false"
  etcd-s3-timeout: 5m
EOF
  chmod 0600 "$configuration_directory/etcd-s3-secret.yaml"

  echo '[k3s] removing the archived legacy single-node installation'
  cluster_ssh "$CLUSTER_NINGBO_HOST" 'bash -s' <<'REMOTE'
set -euo pipefail
if command -v k3s >/dev/null 2>&1 && systemctl is-active --quiet k3s; then
  k3s kubectl delete namespace easyai-auth-center-staging \
    --ignore-not-found=true --wait=true --timeout=120s
fi
if [[ -x /usr/local/bin/k3s-uninstall.sh ]]; then
  /usr/local/bin/k3s-uninstall.sh
fi
REMOTE
  for host in "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST"; do
    cluster_ssh "$host" \
      '[[ ! -x /usr/local/bin/k3s-uninstall.sh ]] || /usr/local/bin/k3s-uninstall.sh'
  done

  node_names=(ningbo hongkong los-angeles)
  node_hosts=("$CLUSTER_NINGBO_HOST" "$CLUSTER_HONGKONG_HOST" "$CLUSTER_LOS_ANGELES_HOST")
  for index in 0 1 2; do
    host=${node_hosts[$index]}
    node_name=${node_names[$index]}
    cluster_ssh "$host" 'install -d -m 0700 /etc/rancher/k3s'
    printf '%s\n' "$K3S_TOKEN" | cluster_ssh "$host" \
      'umask 077; cat >/etc/rancher/k3s/cluster-token; chmod 0600 /etc/rancher/k3s/cluster-token'
    cluster_scp "$configuration_directory/$node_name.yaml" "$host:/etc/rancher/k3s/config.yaml"
    cluster_scp "$configuration_directory/registries.yaml" \
      "$host:/etc/rancher/k3s/registries.yaml"
    cluster_ssh "$host" \
      'chmod 0600 /etc/rancher/k3s/config.yaml /etc/rancher/k3s/registries.yaml'
  done

  for index in 0 1 2; do
    host=${node_hosts[$index]}
    echo "[k3s] installing $k3s_version on $host"
    cluster_scp "$k3s_binary" "$host:/usr/local/bin/k3s"
    cluster_ssh "$host" bash -s -- "$k3s_version" "$k3s_sha256" <<'REMOTE'
set -euo pipefail
version=$1
expected_sha256=$2
chmod 0755 /usr/local/bin/k3s
echo "$expected_sha256  /usr/local/bin/k3s" | sha256sum -c -
install_script=$(mktemp)
trap 'rm -f -- "$install_script"' EXIT
curl -fsSL https://get.k3s.io -o "$install_script"
chmod 0700 "$install_script"
INSTALL_K3S_SKIP_DOWNLOAD=true INSTALL_K3S_VERSION="$version" \
  INSTALL_K3S_EXEC=server "$install_script"
k3s --version | grep -F "$version"
REMOTE
    if [[ $index -eq 0 ]]; then
      for _ in $(seq 1 60); do
        if cluster_ssh "$host" 'k3s kubectl get --raw=/readyz >/dev/null 2>&1'; then
          break
        fi
        sleep 2
      done
      cluster_scp "$configuration_directory/etcd-s3-secret.yaml" \
        "$CLUSTER_NINGBO_HOST:/root/easyai-etcd-s3-secret.yaml"
      cluster_ssh "$CLUSTER_NINGBO_HOST" \
        'k3s kubectl apply -f /root/easyai-etcd-s3-secret.yaml >/dev/null && rm -f /root/easyai-etcd-s3-secret.yaml'
    fi
  done

  for _ in $(seq 1 90); do
    ready_count=$(cluster_ssh "$CLUSTER_NINGBO_HOST" \
      "k3s kubectl get nodes --no-headers 2>/dev/null | awk '\$2 == \"Ready\" { count++ } END { print count + 0 }'")
    [[ $ready_count == 3 ]] && break
    sleep 2
  done
  [[ ${ready_count:-0} == 3 ]] || {
    echo 'three K3s servers did not become Ready' >&2
    return 1
  }

  cluster_ssh "$CLUSTER_NINGBO_HOST" 'bash -s' <<'REMOTE'
set -euo pipefail
kubectl='k3s kubectl'
$kubectl label node easyai-ningbo easyai.io/site=ningbo easyai.io/workload=true easyai.io/database=true --overwrite
$kubectl label node easyai-hongkong easyai.io/site=hongkong easyai.io/workload=true easyai.io/database=true --overwrite
$kubectl label node easyai-los-angeles easyai.io/site=los-angeles easyai.io/workload=false easyai.io/database=false --overwrite
$kubectl taint node easyai-los-angeles easyai.io/witness=true:NoSchedule --overwrite
k3s secrets-encrypt status
k3s etcd-snapshot save --name post-bootstrap
REMOTE

  cluster_scp "$CLUSTER_NINGBO_HOST:/etc/rancher/k3s/k3s.yaml" "$kubeconfig"
  chmod 0600 "$kubeconfig"
  sed -i.bak 's#https://127.0.0.1:6443#https://10.77.0.1:6443#' "$kubeconfig"
  rm -f -- "$kubeconfig.bak"
  assert_private_file "$kubeconfig"
  echo "k3s_bootstrap=PASS version=$k3s_version servers=3 kubeconfig=$kubeconfig"
}

if [[ $action == backup || $action == all ]]; then
  backup_legacy_k3s
fi
if [[ $action == install || $action == all ]]; then
  install_cluster
fi