#!/usr/bin/env bash
# init_full_ubuntu.sh
# Full init script for Ubuntu 22.04 -> 24.04
# - Sysctl tuning (TCP/UDP/orphan sockets)
# - Disable swap
# - NIC offload disable (skip bond/docker/lo/veth/tun/tap/virbr/br/vlan)
# - RPS/RFS per rx-queue
# - Insert SSH keys
# - Install base packages
# - Chrony + timezone
# - limits.conf (nofile 10_000_000)
# - NVMe/SSD tuning (safe)
# - DOES NOT disable firewalld/ufw
set -euo pipefail

echo "=== START init_full_ubuntu.sh ==="

info() { echo ">>> $*"; }
warn() { echo "!!! $*" >&2; }

# -------------------------
# 0. make sure running as root
# -------------------------
if [ "$(id -u)" -ne 0 ]; then
  warn "Run as root (sudo). Exiting."
  exit 1
fi

# -------------------------
# 1. Insert SSH keys
# -------------------------
info "1) Inserting SSH keys into /root/.ssh/authorized_keys"
mkdir -p /root/.ssh
chmod 700 /root/.ssh
AUTH="/root/.ssh/authorized_keys"

cat > /root/.ssh/authorized_keys.tmp <<'KEYS'
ssh-rsa [payload removed: 712 characters] rsa-key-20200326
ssh-rsa [payload removed: 372 characters] root@ht218-88.htnet.vn
ssh-rsa [payload removed: 372 characters] thuannd
ssh-rsa [payload removed: 544 characters] root@proxycp.localdomain
ssh-rsa [payload removed: 372 characters] root@srv203.thuannguyen.net
KEYS

# Merge with existing while removing duplicates, keep existing first
touch "$AUTH"
awk 'FNR==NR{a[$0]=1;next} !a[$0]{print $0}' "$AUTH" /root/.ssh/authorized_keys.tmp > "$AUTH".new || true
# If above fails (awk differences), fallback to simple append+dedup:
if [ ! -s "$AUTH".new ]; then
  cat /root/.ssh/authorized_keys.tmp >> "$AUTH"
  awk '!x[$0]++' "$AUTH" > "${AUTH}.dedup" && mv "${AUTH}.dedup" "$AUTH"
else
  mv "$AUTH".new "$AUTH"
fi
chmod 600 "$AUTH"
rm -f /root/.ssh/authorized_keys.tmp

info "SSH keys installed."

# -------------------------
# 2. apt update + install base packages
# -------------------------
info "2) apt update & install base packages"
export DEBIAN_FRONTEND=noninteractive
apt update -y
apt install -y --no-install-recommends \
  htop iftop net-tools ethtool git curl bwm-ng chrony tcpdump \
  iproute2 lsof iotop rsync atop smartmontools ncdu build-essential

systemctl enable --now chrony

# -------------------------
# 3. Disable swap completely
# -------------------------
info "3) Disable swap"
swapoff -a || true
# comment out swap lines in fstab
cp /etc/fstab /etc/fstab.bak.init_$(date +%s)
sed -i.bak -r 's/^\s*([^#].*\s+swap\s+)/#\1/' /etc/fstab || true

# -------------------------
# 4. limits.conf
# -------------------------
info "4) Apply /etc/security/limits.conf"
cat > /etc/security/limits.conf <<'LIMITS'
# End of file
* soft nofile 10000000
* hard nofile 10000000
root soft nofile 10000000
root hard nofile 10000000
LIMITS

# Try raising current shell limits (may be limited by pam/system)
ulimit -n 10000000 2>/dev/null || warn "ulimit -n could not be raised in this shell."

# -------------------------
# 5. Disable Transparent Huge Pages (THP) persistently
# -------------------------
info "5) Disable Transparent Huge Pages (THP)"
if [ -w /sys/kernel/mm/transparent_hugepage/enabled ]; then
  echo never > /sys/kernel/mm/transparent_hugepage/enabled || true
fi
if [ -w /sys/kernel/mm/transparent_hugepage/defrag ]; then
  echo never > /sys/kernel/mm/transparent_hugepage/defrag || true
fi

cat > /etc/systemd/system/disable-thp.service <<'THP'
[Unit]
Description=Disable Transparent Huge Pages (THP)
After=network.target

[Service]
Type=oneshot
ExecStart=/bin/sh -c 'echo never > /sys/kernel/mm/transparent_hugepage/enabled || true'
ExecStart=/bin/sh -c 'echo never > /sys/kernel/mm/transparent_hugepage/defrag || true'
RemainAfterExit=yes

[Install]
WantedBy=multi-user.target
THP

systemctl daemon-reload
systemctl enable --now disable-thp.service || true

# -------------------------
# 6. Sysctl tuning (TCP/UDP/orphan sockets/etc)
# -------------------------
info "6) Apply sysctl tuning"
cat > /etc/sysctl.conf <<'SYSCTL'
# High-performance kernel tuning for Game/Streaming
fs.file-max = 10000000
fs.nr_open = 10000000

# ports
net.ipv4.ip_local_port_range = 10000 65000

### GIẢM ORPHAN / TIME_WAIT ###
net.ipv4.tcp_fin_timeout = 10
net.ipv4.tcp_tw_reuse = 1
net.ipv4.tcp_orphan_retries = 0
net.ipv4.tcp_max_orphans = 2621440

### BACKLOG ###
net.core.somaxconn = 65535
net.core.netdev_max_backlog = 500000
net.ipv4.tcp_max_syn_backlog = 65535

# Increase the maximum amount of memory buffers
net.core.optmem_max=25165824
# Increase the default and maximum send/receive buffers
net.core.rmem_default=31457280
net.core.rmem_max=67108864
net.core.wmem_default=31457280
net.core.wmem_max=67108864

# UDP tuning
net.ipv4.tcp_mem=786432 1697152 1945728
net.ipv4.udp_mem=65536 131072 262144
net.ipv4.tcp_rmem=4096 87380 33554432
net.ipv4.udp_rmem_min=16384
net.ipv4.tcp_wmem=4096 87380 33554432
net.ipv4.udp_wmem_min=16384
net.ipv4.tcp_max_tw_buckets=1440000
net.ipv4.ip_local_port_range=1024 65535

# other
net.ipv4.tcp_syncookies = 1
net.ipv4.tcp_sack = 1
net.ipv4.tcp_timestamps = 1
net.ipv4.tcp_slow_start_after_idle = 0
net.ipv4.conf.all.rp_filter=1
net.ipv4.tcp_synack_retries=3
net.netfilter.nf_conntrack_max=10000000
net.netfilter.nf_conntrack_tcp_loose=0
net.netfilter.nf_conntrack_tcp_timeout_established=1800
net.netfilter.nf_conntrack_tcp_timeout_close=10
net.netfilter.nf_conntrack_tcp_timeout_close_wait=10
net.netfilter.nf_conntrack_tcp_timeout_fin_wait=20
net.netfilter.nf_conntrack_tcp_timeout_last_ack=20
net.netfilter.nf_conntrack_tcp_timeout_syn_recv=20
net.netfilter.nf_conntrack_tcp_timeout_syn_sent=20
net.netfilter.nf_conntrack_tcp_timeout_time_wait=10

# rps default
net.core.rps_sock_flow_entries = 32768
SYSCTL

# Apply sysctl
sysctl --system || warn "sysctl --system reported errors (check above)"

# -------------------------
# 7. NIC offload & buffer & coalesce tuning (skip virtual types)
# -------------------------
info "7) NIC tuning: disable offloads on active NICs (skipping virtual/bridge/bond/docker/veth/tun/tap)"
SKIP_RE='^(lo$|bond|docker|veth|br-|virbr|tun|tap|vxlan|vlan|cni|macvtap)'
for ifname in $(ls /sys/class/net); do
  if [[ "$ifname" =~ $SKIP_RE ]]; then
    echo "  - skip $ifname (virtual/system)"
    continue
  fi

  # check whether operstate exists and its value
  state="unknown"
  if [ -f "/sys/class/net/$ifname/operstate" ]; then
    state=$(cat /sys/class/net/$ifname/operstate || echo unknown)
  fi

  # Only tune if interface exists (we allow tuning even if down but prefer up)
  if [ "$state" != "down" ]; then
    echo "  - tuning $ifname (state=$state)"
    # ring sizes (may fail silently on some drivers)
    ethtool -G "$ifname" rx 4096 tx 4096 >/dev/null 2>&1 || true
    # reduce coalescing latency (may not be supported)
    ethtool -C "$ifname" rx-usecs 25 rx-frames 64 tx-usecs 25 tx-frames 64 >/dev/null 2>&1 || true
    # disable offloads that often cause issues under high packet rates
    ethtool -K "$ifname" gro off gso off tso off lro off sg off >/dev/null 2>&1 || true
  else
    echo "  - skip $ifname (down)"
  fi
done

# -------------------------
# 8. RPS / RFS (per rx queue)
# -------------------------
info "8) Setup RPS/RFS per rx-queue"
CORES=$(nproc)
# limit mask width to max 64 hex nibbles => 256 bits, but practically limit to 128 bits.
# We'll limit to first 64 cores (most servers). If you have >64 cores and want full spread, tell me.
RPS_CORES=$(( CORES > 64 ? 64 : CORES ))
# number of hex digits needed = ceil(RPS_CORES / 4)
HEX_DIGITS=$(( (RPS_CORES + 3) / 4 ))
# create mask: 'f' repeated HEX_DIGITS times
MASK=$(printf 'f%.0s' $(seq 1 $HEX_DIGITS))
info "  cores=$CORES rps_cores=$RPS_CORES hex_digits=$HEX_DIGITS mask=0x$MASK"

for ifname in $(ls /sys/class/net); do
  if [ ! -d "/sys/class/net/$ifname/queues" ]; then
    continue
  fi
  if [[ "$ifname" =~ $SKIP_RE ]]; then
    continue
  fi
  for q in /sys/class/net/$ifname/queues/rx-*; do
    if [ -f "$q/rps_cpus" ]; then
      printf '%s' "$MASK" > "$q/rps_cpus" 2>/dev/null || true
    fi
    if [ -f "$q/rps_flow_cnt" ]; then
      printf '%s\n' 32768 > "$q/rps_flow_cnt" 2>/dev/null || true
    fi
  done
done

# ensure global rps sock flow entries
if [ -w /proc/sys/net/core/rps_sock_flow_entries ]; then
  echo 32768 > /proc/sys/net/core/rps_sock_flow_entries || true
fi

# -------------------------
# 9. NVMe / SSD tuning (safe)
# -------------------------
info "9) NVMe/SSD tuning (safe settings)"
# operate on all block devices that are 'disk' as returned by lsblk
for disk in $(lsblk -nd -o NAME,TYPE | awk '$2=="disk"{print $1}'); do
  path="/sys/block/${disk}/queue"
  if [ ! -d "$path" ]; then
    continue
  fi
  echo "  - tuning /sys/block/${disk}"

  # try recommended scheduler (mq-deadline) if present, else noop
  if [ -w "$path/scheduler" ]; then
    if grep -q mq-deadline "$path/scheduler" 2>/dev/null; then
      echo mq-deadline > "$path/scheduler" 2>/dev/null || true
    else
      # try 'none' if supported
      echo none > "$path/scheduler" 2>/dev/null || true
    fi
  fi

  # reduce merges
  if [ -w "$path/nomerges" ]; then
    echo 2 > "$path/nomerges" 2>/dev/null || true
  fi

  # rq_affinity might exist
  if [ -w "$path/rq_affinity" ]; then
    echo 2 > "$path/rq_affinity" 2>/dev/null || true
  fi

  # adjust read-ahead
  if [ -w "$path/read_ahead_kb" ]; then
    echo 128 > "$path/read_ahead_kb" 2>/dev/null || true
  fi

  # safe disable of wbt (use 'off' if writable)
  if [ -w "$path/wbt_lat_usec" ]; then
    # many kernels accept "off" rather than numeric 0
    echo off > "$path/wbt_lat_usec" 2>/dev/null || true
  fi

  # nr_requests (if supported)
  if [ -w "$path/nr_requests" ]; then
    echo 1024 > "$path/nr_requests" 2>/dev/null || true
  fi
done

# -------------------------
# 10. small runtime tweaks
# -------------------------
info "10) Misc runtime tweaks"
# file-max
if [ -w /proc/sys/fs/file-max ]; then
  echo 20000000 > /proc/sys/fs/file-max || true
fi

# -------------------------
# 11. reload sysctl to ensure persisted
# -------------------------
info "11) Reloading sysctl settings"
sysctl --system || warn "sysctl reload returned non-zero"

# -------------------------
# 12. final
# -------------------------
info "=== FINISHED init_full_ubuntu.sh ==="
info "Recommended: reboot server to ensure all kernel-level settings (THP disable, schedulers) applied."
echo "Backups: /etc/fstab.bak.init_* created. limits applied to /etc/security/limits.conf."
