#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
PKG_ROOT="$ROOT_DIR"
ENV_FILE="$PKG_ROOT/compose/.env"
COMPOSE_FILE="$PKG_ROOT/compose/docker-compose.yml"

info(){ echo -e "\033[34m[INSTALL-GPU]\033[0m $*"; }
err(){ echo -e "\033[31m[ERROR]\033[0m $*" >&2; }
require(){ local ok=1; for c in "$@"; do command -v "$c" >/dev/null 2>&1 || { err "缺少依赖: $c"; ok=0; }; done; [[ $ok -eq 1 ]]; }
# Compose 检测：优先 docker compose（v2），回退 docker-compose（v1）
require_compose(){
  if docker compose version >/dev/null 2>&1; then return 0; fi
  if command -v docker-compose >/dev/null 2>&1 && docker-compose version >/dev/null 2>&1; then return 0; fi
  err "未检测到 Docker Compose，请安装 docker compose v2 或 docker-compose v1"; exit 1
}
require docker nvidia-smi
require_compose

[[ -f "$ENV_FILE" ]] || { err "缺少 compose/.env，请先运行 scripts/config.sh"; exit 1; }
info "使用环境文件: $ENV_FILE"

# 预热 overlay（当 config 执行很久之前或容器已被清理时，warmup 可能不存在）
set -a; source "$ENV_FILE"; set +a
NET_NAME="${ARGUS_OVERLAY_NET:-argus-sys-net}"
info "检查 overlay 网络可见性: $NET_NAME"
if ! docker network inspect "$NET_NAME" >/dev/null 2>&1; then
  # 如 Overlay 不可见，尝试用 busybox 预热（仅为确保 worker 节点已加入 overlay）
  if ! docker image inspect busybox:latest >/dev/null 2>&1; then
    if [[ -f "$PKG_ROOT/images/busybox.tar" ]]; then docker load -i "$PKG_ROOT/images/busybox.tar"; else err "缺少 busybox 镜像（images/busybox.tar 或本地 busybox:latest）"; exit 1; fi
  fi
  docker rm -f argus-net-warmup >/dev/null 2>&1 || true
  docker run -d --rm --name argus-net-warmup --network "$NET_NAME" busybox:latest sleep 600 >/dev/null 2>&1 || true
  for i in {1..60}; do docker network inspect "$NET_NAME" >/dev/null 2>&1 && break; sleep 1; done
  docker network inspect "$NET_NAME" >/dev/null 2>&1 || { err "预热后仍未看到 overlay: $NET_NAME；请确认 manager 已创建并网络可达"; exit 1; }
  info "overlay 已可见（warmup=argus-net-warmup）"
fi

# 若本函数内重新创建了 warmup 容器，同样测试一次 alias 数据通路
if docker ps --format '{{.Names}}' | grep -q '^argus-net-warmup$'; then
  if ! docker exec argus-net-warmup sh -lc "ping -c 1 -W 2 master.argus.com >/dev/null 2>&1"; then
    err "GPU install 阶段：warmup 容器内无法通过别名访问 master.argus.com；请检查 overlay $NET_NAME 与 server 状态"
    exit 1
  fi
  info "GPU install 阶段：warmup 容器内可达 master.argus.com"
fi

# 导入 GPU bundle 镜像
IMG_TGZ=$(ls -1 "$PKG_ROOT"/images/argus-sys-metric-test-node-bundle-gpu-*.tar.gz 2>/dev/null | head -1 || true)
[[ -n "$IMG_TGZ" ]] || { err "找不到 GPU bundle 镜像 tar.gz"; exit 1; }
info "导入 GPU bundle 镜像: $(basename "$IMG_TGZ")"
tmp=$(mktemp); gunzip -c "$IMG_TGZ" > "$tmp"; docker load -i "$tmp" >/dev/null; rm -f "$tmp"

# 确保日志目录存在（宿主侧，用于映射 /logs/infer 与 /logs/train），并赋权 1777（粘滞位）
mkdir -p "$PKG_ROOT/logs/infer" "$PKG_ROOT/logs/train"
chmod 1777 "$PKG_ROOT/logs/infer" "$PKG_ROOT/logs/train" || true
info "日志目录已准备并赋权 1777: logs/infer logs/train"
stat -c '%a %U:%G %n' "$PKG_ROOT/logs/infer" "$PKG_ROOT/logs/train" 2>/dev/null || true

# 启动 compose 并跟踪日志
PROJECT="${COMPOSE_PROJECT_NAME:-argus-client}"
info "启动 GPU 节点 (docker compose -p $PROJECT up -d)"
docker compose -p "$PROJECT" --env-file "$ENV_FILE" -f "$COMPOSE_FILE" up -d
docker compose -p "$PROJECT" --env-file "$ENV_FILE" -f "$COMPOSE_FILE" ps

# 再次校准宿主日志目录权限，避免容器内脚本对 bind mount 权限回退
chmod 1777 "$PKG_ROOT/logs/infer" "$PKG_ROOT/logs/train" || true
stat -c '%a %U:%G %n' "$PKG_ROOT/logs/infer" "$PKG_ROOT/logs/train" 2>/dev/null || true

info "跟踪节点容器日志（按 Ctrl+C 退出）"
docker logs -f argus-metric-gpu-node-swarm || true