#!/bin/bash
set -euo pipefail

# ────────────────────────────────────────────────
#  颜色输出
# ────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; BLUE='\033[0;34m'; NC='\033[0m'
info()    { echo -e "${BLUE}[INFO]${NC}  $*"; }
success() { echo -e "${GREEN}[OK]${NC}    $*"; }
warn()    { echo -e "${YELLOW}[WARN]${NC}  $*"; }
error()   { echo -e "${RED}[ERROR]${NC} $*"; exit 1; }

# ────────────────────────────────────────────────
#  权限检查
# ────────────────────────────────────────────────
[[ $EUID -ne 0 ]] && error "请用 root 权限运行此脚本: sudo bash $0"

# ────────────────────────────────────────────────
#  欢迎界面
# ────────────────────────────────────────────────
clear
echo "=================================================="
echo "       瓦力方程 WalleCube Smart UPS 一键配置脚本"
echo "       支持型号: W150 / W180"
echo "=================================================="
echo ""

# ────────────────────────────────────────────────
#  收集用户参数（含校验）
# ────────────────────────────────────────────────
# UPS 名称：仅允许字母/数字/下划线/连字符
while true; do
    read -rp "UPS 名称          [默认: ups]:       " UPS_NAME
    UPS_NAME="${UPS_NAME:-ups}"
    if [[ "$UPS_NAME" =~ ^[a-zA-Z0-9_-]+$ ]]; then
        break
    fi
    warn "UPS 名称只允许字母、数字、下划线(_)、连字符(-)，请重新输入。"
done

read -rp "UPS 描述          [默认: W150]:      " UPS_DESC
UPS_DESC="${UPS_DESC:-W150}"
# 描述里不允许双引号，防止破坏 desc = "..."
UPS_DESC="${UPS_DESC//\"/}"

# 监控密码：白名单，只允许常见安全字符，避免 shell 注入与 NUT 解析问题
while true; do
    read -rp "监控用户密码      [默认: mypassword]: " MONITOR_PASS
    MONITOR_PASS="${MONITOR_PASS:-mypassword}"
    if [[ "$MONITOR_PASS" =~ ^[a-zA-Z0-9_!@#%^*+=:.,-]+$ ]]; then
        break
    fi
    warn "密码只允许字母、数字、以及 _ ! @ # % ^ * + = : . , -，请重新输入。"
done

# 关机延迟：必须是非负整数
while true; do
    read -rp "断电后关机延迟(秒)[默认: 30]:        " SHUTDOWN_DELAY
    SHUTDOWN_DELAY="${SHUTDOWN_DELAY:-30}"
    if [[ "$SHUTDOWN_DELAY" =~ ^[0-9]+$ ]]; then
        break
    fi
    warn "关机延迟必须是非负整数（秒），请重新输入。"
done

echo ""
echo "──────────────────────────────────────────────────"
info "配置参数确认:"
echo "  UPS 名称    : $UPS_NAME"
echo "  UPS 描述    : $UPS_DESC"
echo "  监控密码    : $MONITOR_PASS"
echo "  关机延迟    : ${SHUTDOWN_DELAY}s"
echo "──────────────────────────────────────────────────"
read -rp "确认以上配置，开始安装? [Y/n]: " CONFIRM
[[ "${CONFIRM,,}" == "n" ]] && echo "已取消。" && exit 0

echo ""

# ────────────────────────────────────────────────
#  1. 安装 NUT
# ────────────────────────────────────────────────
info "安装 NUT 软件包..."
if command -v apt-get >/dev/null 2>&1; then
    PKG_MANAGER="apt"
    apt-get update -qq
    apt-get install -y -qq nut nut-client nut-server
elif command -v dnf >/dev/null 2>&1; then
    PKG_MANAGER="dnf"
    # 注意：RHEL/Fedora 系没有 nut-server 包，nut 包本身就含服务端
    dnf install -y -q nut nut-client
elif command -v yum >/dev/null 2>&1; then
    PKG_MANAGER="yum"
    yum install -y -q nut nut-client
elif command -v zypper >/dev/null 2>&1; then
    PKG_MANAGER="zypper"
    zypper --non-interactive install nut
elif command -v pacman >/dev/null 2>&1; then
    PKG_MANAGER="pacman"
    pacman -Sy --noconfirm nut
else
    error "未检测到受支持的包管理器（apt/dnf/yum/zypper/pacman），请手动安装 NUT 后再运行。"
fi
success "NUT 安装完成 ($PKG_MANAGER)"

# ────────────────────────────────────────────────
#  2. 探测配置目录、运行时目录、用户/组
# ────────────────────────────────────────────────
# 配置目录：RHEL 7 等旧系统在 /etc/ups，新系统统一在 /etc/nut
if [ -d /etc/nut ]; then
    NUT_DIR=/etc/nut
elif [ -d /etc/ups ]; then
    NUT_DIR=/etc/ups
else
    # 兜底创建 /etc/nut
    mkdir -p /etc/nut
    NUT_DIR=/etc/nut
fi
info "配置目录: $NUT_DIR"

# 用户/组探测
NUT_GROUP="nut"
if ! getent group "$NUT_GROUP" >/dev/null 2>&1; then
    if getent group ups >/dev/null 2>&1; then
        NUT_GROUP="ups"
    else
        NUT_GROUP="root"
    fi
fi

RUNTIME_USER="nut"
if ! id -u "$RUNTIME_USER" >/dev/null 2>&1; then
    if id -u ups >/dev/null 2>&1; then
        RUNTIME_USER="ups"
    else
        RUNTIME_USER="root"
    fi
fi
info "运行身份: ${RUNTIME_USER}:${NUT_GROUP}"

# 探测 NUT 版本，决定使用 master(<2.8) 还是 primary(>=2.8)
NUT_ROLE_KEYWORD="master"
NUT_VER_STR=""
for bin in upsd upsmon upsc; do
    if command -v "$bin" >/dev/null 2>&1; then
        NUT_VER_STR=$("$bin" -V 2>&1 | grep -oE '[0-9]+\.[0-9]+(\.[0-9]+)?' | head -1 || true)
        [ -n "$NUT_VER_STR" ] && break
    fi
done
if [ -n "$NUT_VER_STR" ]; then
    NUT_MAJOR="${NUT_VER_STR%%.*}"
    NUT_REST="${NUT_VER_STR#*.}"
    NUT_MINOR="${NUT_REST%%.*}"
    if [ "$NUT_MAJOR" -gt 2 ] 2>/dev/null || { [ "$NUT_MAJOR" -eq 2 ] 2>/dev/null && [ "$NUT_MINOR" -ge 8 ] 2>/dev/null; }; then
        NUT_ROLE_KEYWORD="primary"
    fi
    info "检测到 NUT 版本: $NUT_VER_STR，使用角色关键字: $NUT_ROLE_KEYWORD"
else
    warn "无法探测 NUT 版本，默认使用旧关键字 'master'（兼容 NUT < 2.8）"
fi

# 运行时目录：优先 /run（systemd 现代写法），不存在则用 /var/run
if [ -d /run ]; then
    NUT_RUN_DIR=/run/nut
else
    NUT_RUN_DIR=/var/run/nut
fi
info "运行时目录: $NUT_RUN_DIR"

# ────────────────────────────────────────────────
#  3. 写入配置文件
# ────────────────────────────────────────────────
info "写入 $NUT_DIR/nut.conf ..."
cat > "$NUT_DIR/nut.conf" <<EOF
MODE=standalone
EOF

info "写入 $NUT_DIR/ups.conf ..."
cat > "$NUT_DIR/ups.conf" <<EOF
[$UPS_NAME]
    driver = usbhid-ups
    port = auto
    desc = "$UPS_DESC"
EOF

info "写入 $NUT_DIR/upsd.conf ..."
cat > "$NUT_DIR/upsd.conf" <<EOF
LISTEN 127.0.0.1 3493
EOF

info "写入 $NUT_DIR/upsd.users ..."
cat > "$NUT_DIR/upsd.users" <<EOF
[upsmon]
    password = $MONITOR_PASS
    upsmon $NUT_ROLE_KEYWORD
EOF

info "写入 $NUT_DIR/upsmon.conf ..."
cat > "$NUT_DIR/upsmon.conf" <<EOF
MONITOR ${UPS_NAME}@localhost 1 upsmon ${MONITOR_PASS} ${NUT_ROLE_KEYWORD}

MINSUPPLIES 1
SHUTDOWNCMD "shutdown -h now"
POWERDOWNFLAG /etc/killpower

NOTIFYCMD /usr/local/bin/upssched-cmd
NOTIFYFLAG ONBATT  SYSLOG+EXEC
NOTIFYFLAG ONLINE  SYSLOG+EXEC
NOTIFYFLAG LOWBATT SYSLOG+EXEC
EOF

info "写入 $NUT_DIR/upssched.conf ..."
cat > "$NUT_DIR/upssched.conf" <<EOF
CMDSCRIPT /usr/local/bin/upssched-cmd
PIPEFN ${NUT_RUN_DIR}/upssched.pipe
LOCKFN ${NUT_RUN_DIR}/upssched.lock

AT ONBATT  * START-TIMER  shutdown-timer ${SHUTDOWN_DELAY}
AT ONLINE  * CANCEL-TIMER shutdown-timer
AT LOWBATT * EXECUTE      shutdown-now
EOF

# ────────────────────────────────────────────────
#  4. 创建 upssched-cmd 脚本
# ────────────────────────────────────────────────
info "创建 /usr/local/bin/upssched-cmd ..."
mkdir -p /usr/local/bin
cat > /usr/local/bin/upssched-cmd <<'SCRIPT'
#!/bin/bash
# upssched-cmd: 由 upssched / upsmon 触发的事件处理脚本
case "$1" in
    shutdown-timer|shutdown-now)
        logger -t upssched-cmd "断电触发，正在执行关机: $1"
        /usr/sbin/shutdown -h now "NUT: 断电，系统关机" 2>/dev/null \
            || /sbin/shutdown -h now "NUT: 断电，系统关机" 2>/dev/null \
            || shutdown -h now "NUT: 断电，系统关机"
        ;;
    *)
        logger -t upssched-cmd "未处理事件: $1"
        ;;
esac
SCRIPT
chmod +x /usr/local/bin/upssched-cmd
success "upssched-cmd 创建完成"

# ────────────────────────────────────────────────
#  5. 设置文件权限与运行时目录（含 tmpfiles 持久化）
# ────────────────────────────────────────────────
info "设置配置文件权限..."
chown "root:${NUT_GROUP}" "$NUT_DIR"/*.conf "$NUT_DIR/upsd.users"
chmod 640 "$NUT_DIR"/*.conf "$NUT_DIR/upsd.users"

# 当前会话立即创建运行时目录
mkdir -p "$NUT_RUN_DIR"
chown "${RUNTIME_USER}:${NUT_GROUP}" "$NUT_RUN_DIR"
chmod 0750 "$NUT_RUN_DIR"

# tmpfiles.d 持久化：重启后 /run/nut 会被 tmpfs 清空，需要 systemd-tmpfiles 自动重建
if [ -d /etc/tmpfiles.d ]; then
    cat > /etc/tmpfiles.d/nut.conf <<EOF
d ${NUT_RUN_DIR} 0750 ${RUNTIME_USER} ${NUT_GROUP} -
EOF
    if command -v systemd-tmpfiles >/dev/null 2>&1; then
        systemd-tmpfiles --create /etc/tmpfiles.d/nut.conf 2>/dev/null || true
    fi
    info "已写入 /etc/tmpfiles.d/nut.conf，重启后自动重建 $NUT_RUN_DIR"
else
    warn "未发现 /etc/tmpfiles.d，重启后 $NUT_RUN_DIR 可能丢失，请自行持久化。"
fi
success "权限设置完成"

# ────────────────────────────────────────────────
#  6. 启动并启用服务
# ────────────────────────────────────────────────
info "启动 NUT 服务..."

# systemd 单元探测：list-unit-files 即使找不到也返回 0，所以必须 grep 实际输出
has_unit() {
    systemctl list-unit-files --no-legend 2>/dev/null \
        | awk '{print $1}' \
        | grep -Fxq "$1"
}

# 先停止所有相关服务（包含驱动模板实例），确保配置生效
systemctl stop 'nut-driver@*' nut-driver nut-driver.target nut-server nut-monitor upsd upsmon 2>/dev/null || true

# 启动 UPS 驱动
if has_unit "nut-driver@.service"; then
    systemctl enable --now "nut-driver@${UPS_NAME}"
    info "已启用 nut-driver@${UPS_NAME}"
elif has_unit "nut-driver.target"; then
    systemctl enable --now nut-driver.target
    info "已启用 nut-driver.target"
elif has_unit "nut-driver.service"; then
    systemctl enable --now nut-driver
    info "已启用 nut-driver.service"
else
    warn "未找到 nut-driver 相关 systemd 单元，尝试直接调用 upsdrvctl..."
    if command -v upsdrvctl >/dev/null 2>&1; then
        upsdrvctl start || warn "upsdrvctl start 失败，请检查 USB 连接。"
    fi
fi

# 等待驱动连接 UPS（USB HID 枚举通常需要 3–5 秒）
sleep 5

# 启动 upsd 服务端
if has_unit "nut-server.service"; then
    systemctl enable --now nut-server
elif has_unit "upsd.service"; then
    systemctl enable --now upsd
else
    warn "未找到 upsd/nut-server 服务单元。"
fi

# 启动 upsmon 监控
if has_unit "nut-monitor.service"; then
    systemctl enable --now nut-monitor
elif has_unit "upsmon.service"; then
    systemctl enable --now upsmon
elif has_unit "nut-client.service"; then
    systemctl enable --now nut-client
else
    warn "未找到 upsmon/nut-monitor 服务单元。"
fi

success "NUT 服务已启动并设为开机自启"

# ────────────────────────────────────────────────
#  7. 验证配置
# ────────────────────────────────────────────────
echo ""
info "验证 NUT 配置..."
sleep 3

echo "──────────────────────────────────────────────────"
if upsc "${UPS_NAME}@localhost" 2>/dev/null | grep -q "device.type\|ups.status"; then
    success "UPS 连接正常！以下为 UPS 状态:"
    echo ""
    upsc "${UPS_NAME}@localhost" 2>/dev/null \
        | grep -E "ups\.(status|load|model)|battery\.(charge|voltage|runtime)|input\.voltage" \
        || true
else
    warn "无法读取 UPS 数据，请检查 USB 连接。"
    warn "排查命令: upsc ${UPS_NAME}@localhost"
    warn "驱动日志: journalctl -u nut-driver@${UPS_NAME} --no-pager -n 50"
fi
echo "──────────────────────────────────────────────────"

echo ""
echo "=================================================="
success "NUT 配置完成！"
echo ""
echo "  常用命令:"
echo "    upsc ${UPS_NAME}@localhost          # 查看 UPS 状态"
echo "    systemctl status nut-server         # 查看服务状态"
echo "    systemctl status nut-monitor        # 查看监控状态"
echo "    journalctl -u nut-driver@${UPS_NAME} # 查看驱动日志"
echo "=================================================="