#!/usr/bin/env bash
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
SKILL_ROOT="$(dirname "$SCRIPT_DIR")"
source "$SCRIPT_DIR/_common.sh"

: "${NODE_NAME:=}"
: "${SOVEREIGN_USER:=sovereign}"
: "${SSH_PUBLIC_KEY:=}"
: "${SSH_PORT:=}"                   # optional; auto-detect if unset
: "${ALLOW_SSH_FALLBACK_22:=true}"  # if SSH_PORT != 22, keep 22 open as safety net
: "${WG_PORT:=51820}"
: "${INSTALL_CLOUDFLARED:=true}"
: "${INSTALL_WIREGUARD:=true}"
: "${OUTPUT_DIR:=$SKILL_ROOT/outputs}"
: "${BACKUP_DIR:=$OUTPUT_DIR/backups}"

install_cloudflared() {
  log_info "Installing cloudflared apt repo + package..."
  need curl
  need gpg

  mkdir -p /usr/share/keyrings
  curl -fsSL https://pkg.cloudflare.com/cloudflare-main.gpg | gpg --dearmor -o /usr/share/keyrings/cloudflare-main.gpg

  local code
  code="$(os_codename)"
  [[ -n "$code" ]] || die "Could not determine OS codename for cloudflared apt repo."

  cat > /etc/apt/sources.list.d/cloudflared.list <<EOF
deb [signed-by=/usr/share/keyrings/cloudflare-main.gpg] https://pkg.cloudflare.com/cloudflared ${code} main
EOF
  apt update
  apt install -y cloudflared
}

detect_sshd_port() {
  # Prefer sshd -T output if available; otherwise fall back to 22.
  local p=""
  if command -v sshd >/dev/null 2>&1; then
    p="$(sshd -T 2>/dev/null | awk '/^port /{print $2; exit}' || true)"
  fi
  echo "${p:-22}"
}

ensure_authorized_key() {
  local user="$1"
  local key_line="$2"
  local homedir="$3"
  local ssh_dir="$homedir/.ssh"
  local auth_file="$ssh_dir/authorized_keys"

  install -d -m 700 -o "$user" -g "$user" "$ssh_dir"
  touch "$auth_file"
  chown "$user:$user" "$auth_file"
  chmod 600 "$auth_file"

  # Normalize comparison to "type base64" (ignore comment).
  local key_norm
  key_norm="$(echo "$key_line" | awk '{print $1" "$2}')"
  if [[ -z "$key_norm" ]]; then
    die "SSH_PUBLIC_KEY appears invalid (expected: <type> <base64> [comment])"
  fi

  if awk '{print $1" "$2}' "$auth_file" | grep -qxF "$key_norm"; then
    log_info "authorized_keys already contains this key (idempotent)."
  else
    echo "$key_line" >> "$auth_file"
    log_info "Appended key to $auth_file"
  fi
}

harden_sshd_config_safely() {
  local cfg="/etc/ssh/sshd_config"
  local before="$BACKUP_DIR/sshd_config.before"

  [[ -f "$cfg" ]] || die "Missing $cfg"

  mkdir -p "$BACKUP_DIR"
  if [[ ! -f "$before" ]]; then
    cp -p "$cfg" "$before"
  fi
  backup_file "$cfg" "$BACKUP_DIR"

  local restore_on_exit="1"
  trap 'if [[ "$restore_on_exit" == "1" ]]; then log_warn "Restoring sshd_config from backup due to failure."; cp -p "$before" "$cfg" || true; reload_ssh_service; fi' EXIT

  # Conservative replacements: add if missing, otherwise modify.
  if grep -qE '^\s*#?\s*PermitRootLogin' "$cfg"; then
    sed -i 's/^\s*#\?\s*PermitRootLogin.*/PermitRootLogin no/' "$cfg"
  else
    echo "PermitRootLogin no" >> "$cfg"
  fi

  if grep -qE '^\s*#?\s*PasswordAuthentication' "$cfg"; then
    sed -i 's/^\s*#\?\s*PasswordAuthentication.*/PasswordAuthentication no/' "$cfg"
  else
    echo "PasswordAuthentication no" >> "$cfg"
  fi

  # Validate before reload.
  sshd -t || die "sshd config validation failed (restoring)."

  restore_on_exit="0"
  trap - EXIT

  reload_ssh_service

  if ! is_ssh_active; then
    die "SSH service is not active after reload. Use rollback/emergency_restore_ssh_ufw.sh from console."
  fi
}

main() {
  require_root
  confirm_gate
  [[ -n "$NODE_NAME" ]] || die "NODE_NAME required"
  [[ -n "$SSH_PUBLIC_KEY" ]] || die "SSH_PUBLIC_KEY required"

  mkdir -p "$OUTPUT_DIR" "$BACKUP_DIR"

  log_info "1) Base packages"
  apt update
  apt install -y gnupg pass git curl jq htop tmux vim ufw fail2ban

  # Ensure OpenSSH server exists (some minimal images may not).
  if ! command -v sshd >/dev/null 2>&1; then
    log_warn "openssh-server not found; installing..."
    apt install -y openssh-server
  fi

  # Ensure host keys exist (sshd -t / sshd -T can fail without them)
  if command -v ssh-keygen >/dev/null 2>&1; then
    ssh-keygen -A >/dev/null 2>&1 || true
  fi

  # Ensure SSH service is enabled/running
  restart_ssh_service

  if [[ "$INSTALL_WIREGUARD" == "true" ]]; then
    apt install -y wireguard
  fi

  if [[ "$INSTALL_CLOUDFLARED" == "true" ]]; then
    install_cloudflared
  fi

  log_info "2) Create sovereign user + authorized_keys (idempotent)"
  if id "$SOVEREIGN_USER" >/dev/null 2>&1; then
    log_warn "User exists: $SOVEREIGN_USER"
  else
    useradd -m -s /bin/bash -G sudo "$SOVEREIGN_USER"
  fi

  homedir="$(getent passwd "$SOVEREIGN_USER" | cut -d: -f6)"
  [[ -n "$homedir" ]] || die "Could not resolve homedir for $SOVEREIGN_USER"
  ensure_authorized_key "$SOVEREIGN_USER" "$SSH_PUBLIC_KEY" "$homedir"

  log_info "3) Set hostname"
  hostnamectl set-hostname "$NODE_NAME"
  if ! grep -q "127.0.1.1 $NODE_NAME" /etc/hosts 2>/dev/null; then
    echo "127.0.1.1 $NODE_NAME" >> /etc/hosts
  fi

  log_info "4) Configure UFW (SSH + WG ports BEFORE enable)"
  backup_file "/etc/ufw/user.rules" "$BACKUP_DIR"

  # Auto-detect current sshd port if SSH_PORT unset.
  current_port="$(detect_sshd_port)"
  if [[ -z "$SSH_PORT" ]]; then
    SSH_PORT="$current_port"
  fi

  # Always allow the chosen SSH_PORT. If you are migrating ports, keep 22 open as a safety net.
  ufw allow "${SSH_PORT}/tcp"
  if [[ "$ALLOW_SSH_FALLBACK_22" == "true" && "$SSH_PORT" != "22" ]]; then
    ufw allow "22/tcp"
  fi
  # If auto-detected differs (rare), keep it open too.
  if [[ "$current_port" != "$SSH_PORT" ]]; then
    ufw allow "${current_port}/tcp"
  fi

  ufw allow "${WG_PORT}/udp"
  ufw default deny incoming
  ufw default allow outgoing
  ufw --force enable
  ufw status verbose > "$OUTPUT_DIR/ufw_status_after.txt" || true

  log_info "5) Harden SSH (reload-safe + restore-on-failure)"
  harden_sshd_config_safely

  log_info "Bootstrap apply complete."
  log_info "NEXT: open a second session as the sovereign user before closing root."
}
main "$@"