#!/usr/bin/env bash
# =========================================================
# VNC 远程桌面 (TigerVNC + 多桌面) - 动态传参版
# 适用系统: Ubuntu / Debian / CentOS / RHEL / Rocky / Alma
# =========================================================

set -euo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
LOG_PREFIX="[VNC 远程桌面]"

log()  { echo -e "\e[32m${LOG_PREFIX} $*\e[0m" ; }
warn() { echo -e "\e[33m${LOG_PREFIX} WARN: $*\e[0m" >&2 ; }
die()  { echo -e "\e[31m${LOG_PREFIX} ERROR: $*\e[0m" >&2 ; exit 1 ; }

# 1. 检查 root 权限
[ "$(id -u)" -eq 0 ] || die "此脚本必须以 root 权限运行，请使用 sudo。"

# 2. 解析命令行参数 (-u, -p, -P)
VNC_USER=""
VNC_PASS=""
VNC_PORT=""

usage() {
    echo "用法: $0 -u <用户名> -p <密码> -P <端口号或桌面编号>"
    echo "示例: $0 -u user1 -p 123456 -P 5901  (对应桌面号 :1)"
    echo "示例: $0 -u user2 -p 654321 -P 2     (对应端口号 5902)"
    exit 1
}

while getopts "u:p:P:h" opt; do
    case $opt in
        u) VNC_USER="$OPTARG" ;;
        p) VNC_PASS="$OPTARG" ;;
        P) VNC_PORT="$OPTARG" ;;
        h) usage ;;
        *) usage ;;
    esac
done

# 校验参数是否为空
if [[ -z "$VNC_USER" || -z "$VNC_PASS" || -z "$VNC_PORT" ]]; then
    die "参数缺失！\n请使用: $0 -u <用户名> -p <密码> -P <端口号/桌面号>\n使用 $0 -h 查看帮助。"
fi

# 智能换算端口与显示编号 (>= 5900 视为端口，反之视为桌面号)
if [[ "$VNC_PORT" -ge 5900 ]]; then
    VNC_DISP=$((VNC_PORT - 5900))
    CURRENT_PORT=$VNC_PORT
else
    VNC_DISP=$VNC_PORT
    CURRENT_PORT=$((5900 + VNC_DISP))
fi

log "目标用户: $VNC_USER"
log "VNC密码: (已隐藏)"
log "桌面编号: :$VNC_DISP (监听端口: $CURRENT_PORT)"
echo "----------------------------------------"

# 3. 检测 Linux 发行版
. /etc/os-release 2>/dev/null || true
DISTRO_ID="${ID:-unknown}"
DISTRO_LIKE="${ID_LIKE:-}"

case "$DISTRO_ID" in
    ubuntu|debian) PKG="apt-get"; PKG_INSTALL="DEBIAN_FRONTEND=noninteractive apt-get install -y" ;;
    centos|rhel|rocky|almalinux|ol) PKG="yum"; PKG_INSTALL="yum install -y" ;;
    fedora) PKG="dnf"; PKG_INSTALL="dnf install -y" ;;
    *)
        if echo "$DISTRO_LIKE" | grep -qiE "rhel|centos|fedora"; then
            PKG="yum"; PKG_INSTALL="yum install -y"
        elif echo "$DISTRO_LIKE" | grep -qi "debian"; then
            PKG="apt-get"; PKG_INSTALL="DEBIAN_FRONTEND=noninteractive apt-get install -y"
        else
            die "不支持的系统发行版: $DISTRO_ID"
        fi
        ;;
esac

log "更新软件包缓存并安装 VNC 桌面环境依赖..."
if [[ "$PKG" == "apt-get" ]]; then
    apt-get update -qq
    DEBIAN_FRONTEND=noninteractive apt-get install -yqq \
      ubuntu-desktop-minimal \
      dbus-x11 tigervnc-standalone-server tigervnc-common tigervnc-xorg-extension \
      xvfb x11vnc x11-utils 2>/dev/null || \
    DEBIAN_FRONTEND=noninteractive apt-get install -yqq \
      gnome-session gnome-shell gnome-terminal nautilus metacity \
      dbus-x11 tigervnc-standalone-server tigervnc-common \
      xvfb x11vnc x11-utils \
      xserver-xorg-video-dummy \
      gnome-shell-extension-desktop-icons-ng gnome-terminal nautilus-extension-gnome-terminal gnome-tweaks
fi

# 4. 用户及目录配置
log "配置用户 '$VNC_USER' 及工作目录..."
id "$VNC_USER" >/dev/null 2>&1 || useradd -m -s /bin/bash "$VNC_USER"

USER_HOME="$(getent passwd "$VNC_USER" | cut -d: -f6)"
CURRENT_SHELL="$(getent passwd "$VNC_USER" | cut -d: -f7)"

if [ "$CURRENT_SHELL" != "/bin/bash" ]; then
    usermod -s /bin/bash "$VNC_USER"
fi

install -d -o "$VNC_USER" -g "$VNC_USER" -m 0750 "$USER_HOME"
install -d -o "$VNC_USER" -g "$VNC_USER" -m 0700 "$USER_HOME/.vnc"

# 设置 VNC 密码
log "生成 VNC 密码..."
printf "%s\n" "$VNC_PASS" | vncpasswd -f > "$USER_HOME/.vnc/passwd"
chmod 600 "$USER_HOME/.vnc/passwd"
chown "$VNC_USER":"$VNC_USER" "$USER_HOME/.vnc/passwd"

# 5. 写入 xstartup (桌面启动脚本)
log "写入 VNC 桌面环境启动脚本 (xstartup)..."
cat > "$USER_HOME/.vnc/xstartup" <<'XSTARTUP_EOF'
#!/bin/sh
unset SESSION_MANAGER
unset DBUS_SESSION_BUS_ADDRESS
export XDG_CURRENT_DESKTOP=ubuntu:GNOME
export XDG_SESSION_TYPE=x11
export GNOME_SHELL_SESSION_MODE=ubuntu
exec dbus-launch --exit-with-session gnome-session
XSTARTUP_EOF
chmod +x "$USER_HOME/.vnc/xstartup"
chown "$VNC_USER":"$VNC_USER" "$USER_HOME/.vnc/xstartup"

# 6. 配置 Systemd 服务
SERVICE_FILE="/etc/systemd/system/vncserver-${VNC_USER}.service"
log "写入专属 systemd 守护进程服务: $SERVICE_FILE"

cat > "$SERVICE_FILE" <<UNIT_EOF
[Unit]
Description=TigerVNC server for $VNC_USER on display :$VNC_DISP (gnome)
After=syslog.target network.target

[Service]
Type=simple
User=$VNC_USER
Group=$VNC_USER
WorkingDirectory=$USER_HOME

# 启动前强制清理残留锁文件，防止 98 端口被占错误
ExecStartPre=-/usr/bin/vncserver -kill :$VNC_DISP
ExecStartPre=-/bin/sh -c 'rm -f /tmp/.X${VNC_DISP}-lock /tmp/.X11-unix/X${VNC_DISP}'

# 启动 VNC 服务 (允许所有 IP 连接 localhost=0)
ExecStart=/usr/bin/vncserver -fg :$VNC_DISP -geometry 1920x1080 -depth 24 -localhost=0
ExecStop=/usr/bin/vncserver -kill :$VNC_DISP

[Install]
WantedBy=multi-user.target
UNIT_EOF

# 7. 端口防冲突及服务启动
if ss -xlpn 2>/dev/null | grep -q "/tmp/.X11-unix/X${VNC_DISP} "; then
    warn "⚠️ 桌面编号 :$VNC_DISP (端口 $CURRENT_PORT) 似乎正被其他进程占用！"
    warn "我们将尝试强制接管该端口..."
fi

log "清理底层的历史残留套接字..."
rm -f /tmp/.X${VNC_DISP}-lock /tmp/.X11-unix/X${VNC_DISP} 2>/dev/null || true
pkill -f "Xvnc.*:${VNC_DISP} " 2>/dev/null || true
pkill -f "Xtigervnc.*:${VNC_DISP}" 2>/dev/null || true

log "启动 VNC 服务..."
systemctl daemon-reload
systemctl enable "vncserver-${VNC_USER}.service"
systemctl restart "vncserver-${VNC_USER}.service"

sleep 3

# 8. 检查结果
if systemctl is-active --quiet "vncserver-${VNC_USER}.service"; then
    log "🎯 部署完毕！"
    echo "====================================================="
    echo "  用户: $VNC_USER"
    echo "  服务: systemctl status vncserver-${VNC_USER}.service"
    echo "  连接: 使用 VNC Viewer 连接 <服务器IP>:$CURRENT_PORT"
    echo "====================================================="
else
    die "服务启动失败，请运行 'journalctl -xeu vncserver-${VNC_USER}.service' 查看日志。"
fi