#!/usr/bin/env bash
# 支持 root 会话的修复版 GNOME + xrdp + TigerVNC 安装脚本
# 模式: full|minimal|core  (core 不依赖 snapd/firefox)
# 新增: ALLOW_ROOT=1 允许使用 root 作为 TARGET_USER
#
# 环境变量:
#   DESKTOP_MODE=full|minimal|core (默认 full)
#   SNAP_POLICY=install|switch-core|fail (默认 install)
#   TARGET_USER=<用户名> (默认自动推断；若最终为 root 需 ALLOW_ROOT=1)
#   ALLOW_ROOT=1 允许 root 作为图形用户
#   VNC_PASS=xxx
#   VNC_GEOMETRY=1920x1080
#   VNC_DEPTH=24
#   AUTO_DISABLE_ANIMATION=1|0 (root 下会自动跳过)
#   HEALTH_CHECK=1|0
#
# 参数:
#   --verify 同 HEALTH_CHECK=1
#   --no-firefox 提示绕过 firefox (需改 core)
#
set -euo pipefail

NO_FIREFOX=0
RUN_VERIFY="${HEALTH_CHECK:-0}"
for a in "$@"; do
  case "$a" in
    --verify) RUN_VERIFY=1 ;;
    --no-firefox) NO_FIREFOX=1 ;;
  esac
done

log() { printf "%s %s\n" "$1" "$2"; }
die() { log "ERROR" "$1"; exit 1; }

if [[ $EUID -ne 0 ]]; then die "请用 root 或 sudo 运行"; fi
command -v apt-get >/dev/null 2>&1 || die "仅支持 apt 的 Ubuntu"

DESKTOP_MODE="${DESKTOP_MODE:-full}"
SNAP_POLICY="${SNAP_POLICY:-install}"
VNC_PASS="${VNC_PASS:-aa123123}"
VNC_GEOMETRY="${VNC_GEOMETRY:-1920x1080}"
VNC_DEPTH="${VNC_DEPTH:-24}"
AUTO_DISABLE_ANIMATION="${AUTO_DISABLE_ANIMATION:-1}"
ALLOW_ROOT="${ALLOW_ROOT:-1}"

# ---------- 用户解析 (支持 root) ----------
if [[ -n "${TARGET_USER:-}" ]]; then
  if ! id "$TARGET_USER" &>/dev/null; then
    die "指定的 TARGET_USER=$TARGET_USER 不存在"
  fi
else
  CAND="${SUDO_USER:-$(logname 2>/dev/null || true)}"
  if [[ -n "$CAND" && "$CAND" != "root" ]] && id "$CAND" &>/dev/null; then
    TARGET_USER="$CAND"
  else
    TARGET_USER="root"
  fi
fi

if [[ "$TARGET_USER" == "root" && "$ALLOW_ROOT" != "1" ]]; then
  die "将使用 root 作为桌面用户。请显式 ALLOW_ROOT=1 以确认: sudo ALLOW_ROOT=1 TARGET_USER=root $0"
fi

HOME_DIR="$(eval echo "~$TARGET_USER")"
[[ -d "$HOME_DIR" ]] || die "用户家目录不存在: $HOME_DIR"

export DEBIAN_FRONTEND=noninteractive

log INFO "================ 配置概要 ================"
log INFO "目标用户:            $TARGET_USER"
log INFO "HOME:                $HOME_DIR"
log INFO "桌面模式:            $DESKTOP_MODE"
log INFO "SNAP 策略:           $SNAP_POLICY"
log INFO "允许 root:           $ALLOW_ROOT"
log INFO "RDP 端口:            3389"
log INFO "VNC 端口(:1):        5901"
log INFO "VNC 分辨率/深度:     $VNC_GEOMETRY / $VNC_DEPTH"
log INFO "关闭动画:            $AUTO_DISABLE_ANIMATION (root 会跳过)"
log INFO "健康检查:            $RUN_VERIFY"
log INFO "no-firefox 标志:     $NO_FIREFOX"
log INFO "==========================================="
sleep 1

# ---------- snapd 需求判定 ----------
need_snap=0
case "$DESKTOP_MODE" in
  full|minimal) need_snap=1 ;;
  core) need_snap=0 ;;
  *) die "DESKTOP_MODE 无效: $DESKTOP_MODE" ;;
esac

snap_available=1
if ! dpkg -s snapd >/dev/null 2>&1; then
  snap_available=0
elif ! systemctl is-active snapd.socket >/dev/null 2>&1 && ! systemctl is-active snapd >/dev/null 2>&1; then
  snap_available=0
fi

snap_pinned_neg=0
if grep -Riq "Package: *snapd" /etc/apt/preferences /etc/apt/preferences.d 2>/dev/null; then
  if grep -R "Package: *snapd" -n /etc/apt/preferences /etc/apt/preferences.d 2>/dev/null \
     | while read -r f; do grep -A2 -i "Package: *snapd" "${f%%:*}"; done \
     | grep -qi "Pin-Priority: *-\?1"; then
    snap_pinned_neg=1
  fi
fi

if [[ $need_snap -eq 1 && $snap_available -eq 0 ]]; then
  case "$SNAP_POLICY" in
    install)
      if [[ $snap_pinned_neg -eq 1 ]]; then
        log WARN "snapd 有负优先级 pin，改用 core 模式"
        DESKTOP_MODE="core"; need_snap=0
      else
        log ACTION "安装 snapd..."
        apt-get update -y
        apt-get install -y snapd
        systemctl enable --now snapd.socket || true
        snap_available=1
      fi
      ;;
    switch-core)
      log ACTION "SNAP_POLICY=switch-core -> 切换 core"
      DESKTOP_MODE="core"; need_snap=0
      ;;
    fail)
      die "需要 snapd 但不可用。改 DESKTOP_MODE=core 或设置 SNAP_POLICY=install"
      ;;
    *) die "SNAP_POLICY 无效: $SNAP_POLICY" ;;
  esac
fi

if [[ $NO_FIREFOX -eq 1 && $DESKTOP_MODE != "core" ]]; then
  log WARN "--no-firefox 仅提示。完整/最小桌面元包仍会拉 firefox；避免请用 core。"
fi

# ---------- 安装 ----------
log INFO "[1/9] apt 更新"
apt-get update -y
apt-get install -y software-properties-common curl wget locales gsettings-desktop-schemas dbus-x11

log INFO "[2/9] 安装桌面 ($DESKTOP_MODE)"
install_core() {
  apt-get install -y --no-install-recommends \
    gnome-shell gnome-session gdm3 gnome-terminal nautilus \
    adwaita-icon-theme mesa-utils x11-xserver-utils policykit-1 fonts-dejavu-core
}
case "$DESKTOP_MODE" in
  full)    apt-get install -y ubuntu-desktop ;;
  minimal) apt-get install -y ubuntu-desktop-minimal ;;
  core)    install_core ;;
esac

log INFO "[3/9] 禁用 Wayland"
if [[ -f /etc/gdm3/custom.conf ]]; then
  sed -i 's/^#\?WaylandEnable=.*/WaylandEnable=false/' /etc/gdm3/custom.conf
else
  echo -e "[daemon]\nWaylandEnable=false" >/etc/gdm3/custom.conf
fi
systemctl restart gdm3 || true

log INFO "[4/9] 安装 xrdp"
apt-get install -y xrdp
adduser "$TARGET_USER" ssl-cert || true

# 允许 root 登录 (sesman.ini)
if grep -q '^AllowRootLogin=' /etc/xrdp/sesman.ini; then
  sed -i 's/^AllowRootLogin=.*/AllowRootLogin=1/' /etc/xrdp/sesman.ini
else
  sed -i '/^\[Security\]/a AllowRootLogin=1' /etc/xrdp/sesman.ini
fi

XRDP_STARTWM="/etc/xrdp/startwm.sh"
if ! grep -q "CUSTOM_GNOME_SESSION" "$XRDP_STARTWM"; then
  cp "$XRDP_STARTWM" "${XRDP_STARTWM}.bak.$(date +%s)"
  sed -i '/test -x \/etc\/X11\/Xsession/,$d' "$XRDP_STARTWM"
  cat >>"$XRDP_STARTWM" <<'EOF'
# CUSTOM_GNOME_SESSION
export GNOME_SHELL_SESSION_MODE=gnome
export XDG_CURRENT_DESKTOP=GNOME
export GNOME_DESKTOP_SESSION_ID=this-is-deprecated
if [ -r /etc/default/locale ]; then
  . /etc/default/locale
  export LANG LANGUAGE LC_ALL
fi
if [ -f "$HOME/.xsession" ]; then
  . "$HOME/.xsession"
fi
exec /usr/bin/gnome-session
EOF
fi

if [[ ! -f "$HOME_DIR/.xsession" ]]; then
  echo "gnome-session" >"$HOME_DIR/.xsession"
  chown "$TARGET_USER":"$TARGET_USER" "$HOME_DIR/.xsession" || true
fi

systemctl enable xrdp
systemctl restart xrdp

log INFO "[5/9] 安装 TigerVNC"
apt-get install -y tigervnc-standalone-server tigervnc-common
VNC_DIR="$HOME_DIR/.vnc"
install -d -m 700 -o "$TARGET_USER" -g "$TARGET_USER" "$VNC_DIR"

cat >"$VNC_DIR/xstartup" <<'EOF'
#!/bin/sh
unset SESSION_MANAGER
unset DBUS_SESSION_BUS_ADDRESS
export GNOME_SHELL_SESSION_MODE=gnome
export XDG_CURRENT_DESKTOP=GNOME
export GNOME_DESKTOP_SESSION_ID=this-is-deprecated
exec gnome-session &
EOF
chown "$TARGET_USER":"$TARGET_USER" "$VNC_DIR/xstartup" || true
chmod +x "$VNC_DIR/xstartup"

if [[ -z "$VNC_PASS" ]]; then
  printf "输入 VNC 密码(6~8位, 回车跳过): "
  read -r -s VNC_IN; echo
  [[ -n "$VNC_IN" ]] && VNC_PASS="$VNC_IN"
fi
if [[ -n "$VNC_PASS" ]]; then
  su - "$TARGET_USER" -c "mkdir -p ~/.vnc && echo '$VNC_PASS' | vncpasswd -f > ~/.vnc/passwd"
  chown "$TARGET_USER":"$TARGET_USER" "$HOME_DIR/.vnc/passwd" || true
  chmod 600 "$HOME_DIR/.vnc/passwd"
fi

log INFO "[6/9] 创建 VNC systemd 服务"
cat >/etc/systemd/system/vnc@.service <<EOF
[Unit]
Description=TigerVNC server for %i (:1)
After=network.target

[Service]
Type=forking
User=%i
PAMName=login
PIDFile=/home/%i/.vnc/%H:%i.pid
ExecStartPre=/bin/sh -c 'install -d -m 700 /home/%i/.vnc'
ExecStart=/usr/bin/vncserver -geometry ${VNC_GEOMETRY} -depth ${VNC_DEPTH} -localhost no :1
ExecStop=/usr/bin/vncserver -kill :1

[Install]
WantedBy=multi-user.target
EOF
systemctl daemon-reload
systemctl enable vnc@"$TARGET_USER"
if [[ -f "$HOME_DIR/.vnc/passwd" ]]; then
  systemctl start vnc@"$TARGET_USER" || log WARN "VNC 启动失败，查 ~/.vnc/*.log"
else
  log WARN "未设置 VNC 密码，稍后: sudo -u $TARGET_USER vncpasswd && systemctl start vnc@$TARGET_USER"
fi

log INFO "[7/9] 性能优化"
if [[ "$AUTO_DISABLE_ANIMATION" == "1" && "$TARGET_USER" != "root" ]]; then
  su - "$TARGET_USER" -c "gsettings set org.gnome.desktop.interface enable-animations false" || true
elif [[ "$AUTO_DISABLE_ANIMATION" == "1" && "$TARGET_USER" == "root" ]]; then
  log WARN "跳过 root 动画关闭 (缺少普通 DBus 桌面会话)"
fi

log INFO "[8/9] 防火墙"
if command -v ufw >/dev/null 2>&1; then
  ufw allow 3389/tcp || true
  ufw allow 5901/tcp || true
fi

log INFO "[9/9] dpkg 自检"
dpkg --configure -a || true
apt-get -f install -y || true

health_check() {
  echo
  echo ">>> 健康检查"
  local fail=0
  CHK(){ if eval "$2"; then echo "[OK] $1"; else echo "[FAIL] $1"; fail=$((fail+1)); fi; }
  for p in gnome-shell gnome-session gdm3 xrdp tigervnc-standalone-server; do
    CHK "包安装: $p" "dpkg -s $p >/dev/null 2>&1"
  done
  CHK "gdm3 活跃" "systemctl is-active gdm3 >/dev/null 2>&1"
  CHK "xrdp 活跃" "systemctl is-active xrdp >/dev/null 2>&1"
  systemctl is-enabled vnc@"$TARGET_USER" >/dev/null 2>&1 && echo "[OK] VNC 服务启用" || echo "[WARN] VNC 服务未启用"
  systemctl is-active vnc@"$TARGET_USER" >/dev/null 2>&1 && echo "[OK] VNC 运行中" || echo "[WARN] VNC 未运行"
  CHK "RDP 3389 监听" "ss -tln | awk '{print \$4}' | grep -E '(:|\\.)3389\$' >/dev/null"
  CHK "VNC 5901 监听" "ss -tln | awk '{print \$4}' | grep -E '(:|\\.)5901\$' >/dev/null"
  CHK "Wayland 禁用" "grep -q '^WaylandEnable=false' /etc/gdm3/custom.conf"
  [[ $fail -eq 0 ]] && echo "健康检查：全部通过" || echo "健康检查：$fail 项失败"
  echo ">>> 结束"
}

echo
echo "================ 安装完成 ================="
echo "模式: $DESKTOP_MODE"
echo "用户: $TARGET_USER"
[[ "$TARGET_USER" == "root" ]] && echo "警告: 当前使用 root 图形会话，建议尽快创建普通用户。"
echo "RDP: 3389 (选择 Xorg 会话)"
echo "VNC: 5901 (:1 VNC 密码 "$VNC_PASS")"
[[ -z "$VNC_PASS" ]] && echo "未设 VNC 密码: sudo -u $TARGET_USER vncpasswd && systemctl start vnc@$TARGET_USER"
echo "=========================================="

[[ $RUN_VERIFY -eq 1 ]] && health_check

### 如果出现这个错误就运行两次安装脚本
Errors were encountered while processing: /tmp/apt-dpkg-install-OElkaL/0027-firefox_1%3a1snap1-0ubuntu2_amd64.deb needrestart is being skipped since dpkg has failed E: Sub-process /usr/bin/dpkg returned an error code (1)

# 更新系统
sudo apt update
sudo apt -y upgrade

# 安装常用工具
sudo apt install -y curl wget vim htop net-tools unzip tar git

# 设置时区（示例：Asia/Shanghai）
sudo timedatectl set-timezone Asia/Shanghai

# 确认当前用户（避免直接用 root）
whoami

sudo apt install -y xfce4 xfce4-goodies

sudo apt install -y ubuntu-desktop-minimal

sudo apt install -y xrdp
sudo systemctl enable --now xrdp

echo "xfce4-session" | sudo tee /etc/skel/.Xsession
echo "xfce4-session" > ~/.Xsession

echo "startxfce4" > ~/.xsession
chmod +x ~/.xsession

sudo systemctl restart gdm3
sudo systemctl restart xrdp

sudo adduser xrdp ssl-cert
sudo systemctl restart xrdp

sudo ufw allow from <你的办公出口IP>/32 to any port 3389
# 若没有固定 IP（不推荐），临时：
# sudo ufw allow 3389/tcp
sudo ufw status

sudo apt install -y tigervnc-standalone-server tigervnc-common

vncpasswd

vncserver :1
vncserver -kill :1

#!/bin/sh
unset SESSION_MANAGER
unset DBUS_SESSION_BUS_ADDRESS
startxfce4 &

chmod +x ~/.vnc/xstartup

sudo systemctl daemon-reload
sudo systemctl enable --now vnc@$(whoami).service

ssh -L 5901:localhost:5901 user@server_ip

sudo apt install -y x2goserver x2goserver-xsession

journalctl -u xrdp -f
journalctl -u xrdp-sesman -f
tail -f /var/log/xrdp.log /var/log/xrdp-sesman.log
tail -f ~/.vnc/*.log

# 更新
sudo apt update && sudo apt -y upgrade

# 安装轻量桌面 + XRDP
sudo apt install -y xfce4 xfce4-goodies xrdp
echo "startxfce4" > ~/.xsession
sudo adduser $USER ssl-cert
sudo systemctl enable --now xrdp

# 防火墙限制来源
sudo ufw allow from 1.2.3.4 to any port 3389

# VNC (TigerVNC) + SSH 隧道
sudo apt install -y tigervnc-standalone-server
vncpasswd
vncserver :1
vncserver -kill :1
# 编辑 ~/.vnc/xstartup -> startxfce4
vncserver :1 -localhost yes
ssh -L 5901:localhost:5901 user@server_ip  # 本地连接 localhost:5901

# X2Go
sudo apt install -y x2goserver x2goserver-xsession xfce4

# 卸载 XRDP（可选）
sudo systemctl disable --now xrdp
sudo apt purge -y xrdp