#!/bin/bash
# Dazhe Tongyou — dt-node 一键安装脚本
# 用法: curl -sfL https://soft.dzty.club/install.sh | bash -s -- --key <YOUR_API_KEY>
set -e

# ── 颜色 ──
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
info()  { echo -e "${GREEN}[dt]${NC} $1"; }
warn()  { echo -e "${YELLOW}[dt]${NC} $1"; }
err()   { echo -e "${RED}[dt]${NC} $1" >&2; }

# ── 参数 ──
API_KEY=""
API_ADDR="api.dzty.club:443"
GPU="auto"
VRAM="auto"
PRICE="0.5"
TIER="L3"

while [ $# -gt 0 ]; do
    case "$1" in
        --key)       API_KEY="$2"; shift 2 ;;
        --api)       API_ADDR="$2"; shift 2 ;;
        --gpu)       GPU="$2"; shift 2 ;;
        --price)     PRICE="$2"; shift 2 ;;
        --tier)      TIER="$2"; shift 2 ;;
        --help|-h)   echo "Usage: curl -sfL https://get.dazhe.io/install.sh | bash -s -- --key <API_KEY>"; exit 0 ;;
        *)           err "unknown arg: $1"; exit 1 ;;
    esac
done

echo ""
echo "  ╔══════════════════════════════╗"
echo "  ║   Dazhe Tongyou — dt-node   ║"
echo "  ║   GPU Compute Node Agent    ║"
echo "  ╚══════════════════════════════╝"
echo ""

# ── 检查 API Key ──
if [ -z "$API_KEY" ]; then
    echo -n "Enter your Dazhe Tongyou API Key: "
    read -r API_KEY
    if [ -z "$API_KEY" ]; then
        err "API Key is required. Get one at https://api.dazhe.io/v1/accounts"
        exit 1
    fi
fi

# ── 检测架构 ──
ARCH=$(uname -m)
case "$ARCH" in
    x86_64|amd64)  BIN="dt-node-linux-amd64" ;;
    aarch64|arm64) BIN="dt-node-linux-arm64" ;;
    *) err "unsupported architecture: $ARCH (only x86_64 and arm64)"; exit 1 ;;
esac
info "detected architecture: $ARCH"

# ── 检测 GPU ──
if [ "$GPU" = "auto" ]; then
    if command -v nvidia-smi &>/dev/null; then
        GPU=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | head -1 | xargs)
        VRAM=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>/dev/null | head -1 | xargs)
        info "detected GPU: $GPU ($VRAM GB)"
    else
        GPU="unknown"
        VRAM=0
        warn "no NVIDIA GPU detected — running in CPU-only mode"
    fi
fi

# ── 安装路径 ──
BIN_PATH="/usr/local/bin/dt-node"
SERVICE_PATH="/etc/systemd/system/dt-node.service"

# ── 下载二进制 ──
S3_URL="https://soft.dzty.club/dt-node"
info "downloading dt-node for $ARCH..."
if command -v curl &>/dev/null; then
    curl -sL --connect-timeout 15 "$S3_URL" -o /tmp/dt-node 2>/dev/null
elif command -v wget &>/dev/null; then
    wget -q --timeout=15 "$S3_URL" -O /tmp/dt-node 2>/dev/null
else
    err "neither curl nor wget found"
    exit 1
fi

chmod +x /tmp/dt-node
SIZE=$(stat -c%s /tmp/dt-node 2>/dev/null || stat -f%z /tmp/dt-node 2>/dev/null || echo 0)
if [ "$SIZE" -lt 1000000 ]; then
    err "download failed: file is ${SIZE} bytes (expected >1MB)"
    exit 1
fi
info "downloaded $(numfmt --to=iec $SIZE 2>/dev/null || echo $SIZE bytes)"

# ── 安装 ──
mv /tmp/dt-node "$BIN_PATH"
info "installed to $BIN_PATH"

# ── 创建 systemd 服务 ──
cat > "$SERVICE_PATH" << EOF
[Unit]
Description=Dazhe Tongyou - Compute Node Agent
After=network.target

[Service]
Type=simple
ExecStart=$BIN_PATH --api $API_ADDR --token $API_KEY --gpu '$GPU' --vram $VRAM --price $PRICE
Restart=on-failure
RestartSec=5
User=root

[Install]
WantedBy=multi-user.target
EOF

systemctl daemon-reload
systemctl enable dt-node 2>/dev/null
systemctl start dt-node

# ── 验证 ──
sleep 3
if systemctl is-active dt-node >/dev/null 2>&1; then
    info "✅ dt-node is running!"
    info "    node registered with GPU: $GPU ($VRAM GB)"
    info "    gateway: $API_ADDR"
    echo ""
    echo "  Commands:"
    echo "    sudo systemctl status dt-node"
    echo "    sudo journalctl -u dt-node -f"
    echo "    sudo systemctl stop dt-node"
else
    warn "service started but may not be connected yet"
    warn "check logs: journalctl -u dt-node -f"
fi

# ── 清理 ──
rm -f /tmp/dt-node
echo ""
info "installation complete"