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

CNI_VERSION="v1.9.0"
ARCH="amd64"
GVISOR_ARCH="x86_64"
SHIM_PATCHED_VERSION="patched"
SHIM_PATCHED_URL_BASE="https://github.com/swualabs/gvisor-shim-patched/releases/download"
SHIM_INSTALL_PATH="/usr/bin/containerd-shim-runsc-v1"
SBX_CNI_CONF_DIR="/etc/cni/sandboxd.d"
SBX_CNI_CONF_FILE="${SBX_CNI_CONF_DIR}/20-sbxnet.conflist"

log() { echo "[install] $*"; }
need() { command -v "$1" >/dev/null 2>&1 || { echo "missing command: $1"; exit 1; }; }
die() { echo "[install] ERROR: $*" >&2; exit 1; }

preflight_checks() {
  need sudo
  need systemctl
  need curl
  need tar
  need jq
}

enforce_supported_platform() {
  local arch distro
  arch="$(uname -m)"
  if [[ "${arch}" != "x86_64" && "${arch}" != "amd64" ]]; then
    die "this installer supports x86_64 only (got: ${arch})"
  fi

  if [[ ! -r /etc/os-release ]]; then
    die "cannot detect distro (/etc/os-release missing); Ubuntu x86_64 only"
  fi
  # shellcheck source=/dev/null
  . /etc/os-release
  distro="${ID:-}"
  if [[ "${distro}" != "ubuntu" ]]; then
    die "this installer supports Ubuntu x86_64 only (got distro: ${distro:-unknown})"
  fi

  log "Platform check passed: Ubuntu x86_64"
}

detect_arch() {
  local m
  m="$(uname -m)"
  case "${m}" in
    x86_64|amd64)
      ARCH="amd64"
      GVISOR_ARCH="x86_64"
      ;;
    aarch64|arm64)
      ARCH="arm64"
      GVISOR_ARCH="aarch64"
      ;;
    *)
      die "unsupported architecture: ${m}"
      ;;
  esac
  log "Detected arch: cni=${ARCH}, gvisor=${GVISOR_ARCH}"
}

install_containerd_pkg() {
  if command -v containerd >/dev/null 2>&1; then
    return
  fi
  log "Installing containerd package"
  if sudo apt-cache show containerd.io >/dev/null 2>&1; then
    sudo apt install -y containerd.io
    return
  fi
  if sudo apt-cache show containerd >/dev/null 2>&1; then
    sudo apt install -y containerd
    return
  fi
  die "neither containerd.io nor containerd package is available in apt sources"
}

install_base() {
  log "Installing base packages"
  sudo apt update
  sudo apt install -y iproute2 iptables curl jq tar ca-certificates gnupg apparmor apparmor-utils
  install_containerd_pkg

  log "Ensuring containerd service is running"
  sudo systemctl enable --now containerd
  sudo systemctl is-active containerd >/dev/null || die "containerd failed to start"
}

install_cni() {
  log "Installing CNI plugins ${CNI_VERSION}"
  sudo mkdir -p /opt/cni/bin
  curl -fsSL "https://github.com/containernetworking/plugins/releases/download/${CNI_VERSION}/cni-plugins-linux-${ARCH}-${CNI_VERSION}.tgz" | sudo tar -C /opt/cni/bin -xz

  log "Enabling bridge netfilter"
  sudo modprobe br_netfilter
  cat <<'CONF' | sudo tee /etc/sysctl.d/99-sandboxd.conf >/dev/null
net.bridge.bridge-nf-call-iptables = 1
net.bridge.bridge-nf-call-ip6tables = 1
CONF
  sudo sysctl --system >/dev/null
}

install_runsc() {
  log "Installing gVisor runsc"

  download_gvisor_asset runsc /tmp/runsc
  download_patched_shim /tmp/containerd-shim-runsc-v1

  sudo install -m 0755 /tmp/runsc /usr/local/bin/runsc
  sudo install -m 0755 /tmp/containerd-shim-runsc-v1 "${SHIM_INSTALL_PATH}"
  # Keep /usr/local/bin path consistent to avoid mixed shim versions on PATH.
  sudo ln -sfn "${SHIM_INSTALL_PATH}" /usr/local/bin/containerd-shim-runsc-v1
}

download_gvisor_asset() {
  local asset="$1"
  local out="$2"
  local urls=(
    "https://storage.googleapis.com/gvisor/releases/release/latest/${GVISOR_ARCH}/${asset}"
  )

  for u in "${urls[@]}"; do
    if curl -fsSL -o "${out}" "${u}"; then
      log "Downloaded ${asset} from ${u}"
      return 0
    fi
  done

  die "failed to download ${asset}; tried known gVisor release URLs"
}

download_patched_shim() {
  local out="$1"
  local url="${SHIM_PATCHED_URL_BASE}/${SHIM_PATCHED_VERSION}/containerd-shim-runsc-v1"
  if curl -fsSL -o "${out}" "${url}"; then
    log "Downloaded patched shim from ${url} (version=${SHIM_PATCHED_VERSION})"
    return 0
  fi
  die "failed to download patched shim from ${url}"
}

configure_containerd_for_runsc() {
  log "Writing /etc/containerd/config.toml (runsc runtime)"
  sudo mkdir -p /etc/containerd
  cat <<'TOML' | sudo tee /etc/containerd/config.toml >/dev/null
version = 2

[plugins."io.containerd.grpc.v1.cri".containerd]
  snapshotter = "overlayfs"
  default_runtime_name = "runsc"

[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.runsc]
  runtime_type = "io.containerd.runsc.v1"
  pod_annotations = ["dev.gvisor.*"]

[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.runsc.options]
  TypeUrl = "io.containerd.runsc.v1.options"
  ConfigPath = "/etc/containerd/runsc.toml"

[plugins."io.containerd.grpc.v1.cri".containerd.runtimes.runc]
  runtime_type = "io.containerd.runc.v2"

[plugins."io.containerd.grpc.v1.cri".cni]
  bin_dir = "/opt/cni/bin"
  conf_dir = "/etc/cni/sandboxd.d"
  max_conf_num = 1
TOML

  cat <<'RUNSC' | sudo tee /etc/containerd/runsc.toml >/dev/null
[runsc_config]
allow-flag-override = "true"
RUNSC

  sudo systemctl restart containerd
  sudo systemctl is-active containerd >/dev/null || die "containerd restart failed after runsc config"
}

configure_network() {
  log "Writing CNI config (sbxnet)"
  sudo mkdir -p "${SBX_CNI_CONF_DIR}" /var/lib/cni/sbxnet
  if [[ -f /etc/cni/net.d/20-sbxnet.conflist && ! -f "${SBX_CNI_CONF_FILE}" ]]; then
    log "Migrating existing CNI conflist from /etc/cni/net.d to ${SBX_CNI_CONF_DIR}"
    sudo cp /etc/cni/net.d/20-sbxnet.conflist "${SBX_CNI_CONF_FILE}"
  fi
  cat <<'JSON' | sudo tee "${SBX_CNI_CONF_FILE}" >/dev/null
{
  "cniVersion": "1.0.0",
  "name": "sbxnet",
  "plugins": [
    {
      "type": "bridge",
      "bridge": "sbx-br0",
      "isGateway": true,
      "ipMasq": true,
      "hairpinMode": true,
      "ipam": {
        "type": "host-local",
        "ranges": [[{"subnet": "10.89.0.0/16", "gateway": "10.89.0.1"}]],
        "routes": [{"dst": "0.0.0.0/0"}],
        "dataDir": "/var/lib/cni/sbxnet"
      }
    },
    {"type": "loopback"}
  ]
}
JSON

  log "Ensuring iptables global chains"
  sudo iptables -N SANDBOX-FWD 2>/dev/null || true
  sudo iptables -C FORWARD -j SANDBOX-FWD 2>/dev/null || sudo iptables -I FORWARD 1 -j SANDBOX-FWD
  sudo iptables -N SANDBOX-IN 2>/dev/null || true
  sudo iptables -C INPUT -j SANDBOX-IN 2>/dev/null || sudo iptables -I INPUT 1 -j SANDBOX-IN
}

write_runtime_config() {
  log "Writing runtime settings to /var/lib/sandboxd/sbxlet_config.json"
  local runtime_addr="/run/containerd/containerd.sock"
  local config_path="/var/lib/sandboxd/sbxlet_config.json"
  local tmp_path

  sudo mkdir -p /var/lib/sandboxd
  tmp_path="$(mktemp)"

  if [[ -f "${config_path}" ]]; then
    jq \
      --arg runtime_addr "${runtime_addr}" \
      --arg cni_conf_path "${SBX_CNI_CONF_FILE}" \
      '.containerd_address = $runtime_addr
      | .cni_conf_path = $cni_conf_path
      | .default_ephemeral_storage = "128Mi"
      | .rootfs_ratio_percent = 80
      | .tmpfs_ratio_percent = 20' \
      "${config_path}" > "${tmp_path}"
  else
    cat > "${tmp_path}" <<EOF
{
  "containerd_address": "${runtime_addr}",
  "cni_conf_path": "${SBX_CNI_CONF_FILE}",
  "default_ephemeral_storage": "128Mi",
  "rootfs_ratio_percent": 80,
  "tmpfs_ratio_percent": 20
}
EOF
  fi

  sudo install -m 0644 "${tmp_path}" "${config_path}"
  rm -f "${tmp_path}"
}

main() {
  preflight_checks
  enforce_supported_platform
  detect_arch
  install_base
  install_cni
  install_runsc
  configure_containerd_for_runsc
  configure_network
  write_runtime_config

  log "Install complete"
  containerd --version
  ctr version
  runsc --version | head -n 1
  command -v containerd-shim-runsc-v1 >/dev/null
  log "Runtime: gVisor(runsc)"
}

main "$@"