#!/usr/bin/env bash
#
# setup_traffic_guard.sh — 安装/卸载 traffic-guard（月度出站流量预算看门狗）的 systemd timer
#
# 为什么是独立安装器而不是塞进 setup_cdn_bypass_tproxy.sh：
#   * 生命周期独立：guard 保护的是 xray/caddy/hysteria 整台机器的免费额度，没装 cdn-bypass 的 GCP 机器
#     同样需要；卸载 cdn-bypass 也不该把 guard 一起拆掉。
#   * cdn-bypass 脚本的前置动作（解析 vless 链接、探测 xray、nft/策略路由）与 guard 毫无关系，
#     硬塞会让 ACTION 分发和 DRY_RUN 渲染互相牵扯。
#
# 做的事（幂等）：
#   1. 把同目录的 traffic-guard.sh 装到 /usr/local/sbin/traffic-guard.sh
#   2. 写 /etc/traffic-guard.env（GUARD_* 配置，无机密）
#   3. 写 traffic-guard.service（oneshot）+ traffic-guard.timer，enable --now timer，并立刻跑一次 check 建基线
#
# 用法：
#   sudo ACTION=guard-install bash setup_traffic_guard.sh          # 默认动作
#   sudo ACTION=guard-uninstall bash setup_traffic_guard.sh
#   DRY_RUN=1 bash setup_traffic_guard.sh                          # 只渲染到 ./dry-run-output/
#
# 环境变量（默认值）：GUARD_LIMIT_GIB(200) GUARD_STOP_AT_GIB(185)
#   GUARD_SERVICES("xray caddy hysteria-server") GUARD_INTERVAL(5min) GUARD_TZ(America/Los_Angeles)
#   GUARD_EXCLUDE_IFACES("")  GUARD_SRC（traffic-guard.sh 的路径，默认脚本同目录）
#
# 查看：sudo traffic-guard.sh status    日志：journalctl -u traffic-guard.service

set -euo pipefail

ACTION="${ACTION:-guard-install}"
GUARD_LIMIT_GIB="${GUARD_LIMIT_GIB:-200}"
GUARD_STOP_AT_GIB="${GUARD_STOP_AT_GIB:-185}"
GUARD_SERVICES="${GUARD_SERVICES:-xray caddy hysteria-server}"
GUARD_INTERVAL="${GUARD_INTERVAL:-5min}"
GUARD_TZ="${GUARD_TZ:-America/Los_Angeles}"
GUARD_EXCLUDE_IFACES="${GUARD_EXCLUDE_IFACES:-}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
GUARD_SRC="${GUARD_SRC:-$SCRIPT_DIR/traffic-guard.sh}"
GUARD_BIN="${GUARD_BIN:-/usr/local/sbin/traffic-guard.sh}"
GUARD_ENV_FILE="${GUARD_ENV_FILE:-/etc/traffic-guard.env}"
GUARD_STATE_DIR="${GUARD_STATE_DIR:-/var/lib/traffic-guard}"
UNIT_DIR="${UNIT_DIR:-/etc/systemd/system}"
DRY_RUN="${DRY_RUN:-}"

info() { printf '\033[32m[+]\033[0m %s\n' "$*"; }
warn() { printf '\033[33m[!]\033[0m %s\n' "$*" >&2; }
die()  { printf '\033[31m[x]\033[0m %s\n' "$*" >&2; exit 1; }

if [ -n "$DRY_RUN" ]; then
  DRY_RUN_DIR="./dry-run-output"; mkdir -p "$DRY_RUN_DIR"; info "DRY_RUN：只渲染到 $DRY_RUN_DIR"
elif [ "$(id -u)" -ne 0 ]; then
  die "需要 root 运行"
fi

[[ "$GUARD_INTERVAL" =~ ^[0-9]+(s|sec|m|min|h|hr)?$ ]] || die "GUARD_INTERVAL 格式不对（例：5min / 300s / 1h）：$GUARD_INTERVAL"

render_env() {
  cat <<EOF
# traffic-guard 配置（由 setup_traffic_guard.sh 生成，重复运行会整体覆盖；无机密）
GUARD_LIMIT_GIB="$GUARD_LIMIT_GIB"
GUARD_STOP_AT_GIB="$GUARD_STOP_AT_GIB"
GUARD_SERVICES="$GUARD_SERVICES"
GUARD_TZ="$GUARD_TZ"
GUARD_STATE_DIR="$GUARD_STATE_DIR"
GUARD_EXCLUDE_IFACES="$GUARD_EXCLUDE_IFACES"
EOF
}

render_service() {
  cat <<EOF
[Unit]
Description=Traffic guard: stop proxy services before the monthly egress budget is exceeded
Documentation=file:$GUARD_BIN

[Service]
Type=oneshot
EnvironmentFile=$GUARD_ENV_FILE
ExecStart=$GUARD_BIN check
# 只停 GUARD_SERVICES 里的服务，绝不关机（见 traffic-guard.sh 头注释）
EOF
}

render_timer() {
  cat <<EOF
[Unit]
Description=Run traffic guard every $GUARD_INTERVAL

[Timer]
# 开机后 2 分钟先查一次（重启会把被 guard 停掉的服务重新拉起来，要尽快再停回去）
OnBootSec=2min
OnUnitActiveSec=$GUARD_INTERVAL
AccuracySec=30s

[Install]
WantedBy=timers.target
EOF
}

do_install() {
  [ -f "$GUARD_SRC" ] || die "找不到 traffic-guard.sh：$GUARD_SRC（可用 GUARD_SRC 指定）"
  # 先用同一组环境变量让 guard 自己校验一遍（数字/阈值关系/服务名），不合法就在改系统前中止
  GUARD_STATE_DIR="$(mktemp -d)" GUARD_SYSFS_NET=/nonexistent GUARD_LIMIT_GIB="$GUARD_LIMIT_GIB" \
    GUARD_STOP_AT_GIB="$GUARD_STOP_AT_GIB" GUARD_SERVICES="$GUARD_SERVICES" GUARD_TZ="$GUARD_TZ" \
    bash "$GUARD_SRC" status > /dev/null || die "GUARD_* 参数校验失败（见上方报错）"

  if [ -n "$DRY_RUN" ]; then
    render_env > "$DRY_RUN_DIR/traffic-guard.env"
    render_service > "$DRY_RUN_DIR/traffic-guard.service"
    render_timer > "$DRY_RUN_DIR/traffic-guard.timer"
    cp "$GUARD_SRC" "$DRY_RUN_DIR/traffic-guard.sh"
    info "已渲染：traffic-guard.env / traffic-guard.service / traffic-guard.timer / traffic-guard.sh"
    return 0
  fi

  install -m 0755 "$GUARD_SRC" "$GUARD_BIN"
  info "已安装 $GUARD_BIN"
  mkdir -p "$GUARD_STATE_DIR"
  render_env > "$GUARD_ENV_FILE"; chmod 644 "$GUARD_ENV_FILE"
  render_service > "$UNIT_DIR/traffic-guard.service"
  render_timer > "$UNIT_DIR/traffic-guard.timer"
  systemctl daemon-reload
  systemctl enable --now traffic-guard.timer
  info "timer 已启用（每 $GUARD_INTERVAL；开机后 2min 首查）"
  # 立刻建本月基线（首次运行从 0 起算）
  systemctl start traffic-guard.service || warn "首次 check 失败：journalctl -u traffic-guard.service"
  "$GUARD_BIN" status || true
  warn "首次安装的本月基线从 0 开始；如本月已有出站用量，请照 GCP 账单校准：sudo $GUARD_BIN set-used <GiB>"
  info "完成。卸载：sudo ACTION=guard-uninstall bash $0"
}

do_uninstall() {
  if [ -n "$DRY_RUN" ]; then info "DRY_RUN：卸载只列动作"; return 0; fi
  systemctl disable --now traffic-guard.timer 2>/dev/null || true
  systemctl stop traffic-guard.service 2>/dev/null || true
  local was_stopped=0
  [ -f "$GUARD_STATE_DIR/STOPPED" ] && was_stopped=1
  rm -f "$UNIT_DIR/traffic-guard.service" "$UNIT_DIR/traffic-guard.timer" "$GUARD_ENV_FILE"
  systemctl daemon-reload
  if [ "$was_stopped" = 1 ]; then
    # 不自动拉起服务：此刻很可能就是超预算停机状态，由人决定
    warn "traffic-guard 此前已停止了一些服务（标记 $GUARD_STATE_DIR/STOPPED 仍在）；卸载不会自动启动它们。"
    warn "如需恢复：先 sudo $GUARD_BIN resume，再删除 $GUARD_BIN；状态目录已保留。"
  else
    rm -f "$GUARD_BIN"
    rm -rf "$GUARD_STATE_DIR"
  fi
  info "卸载完成"
}

case "$ACTION" in
  guard-install)   do_install;;
  guard-uninstall) do_uninstall;;
  *) die "ACTION 只支持 guard-install / guard-uninstall";;
esac