#!/bin/bash

# 脚本部署工具 - 下载并配置GPU监控脚本
set -euo pipefail
IFS=$'\n\t'

# 配置信息
SCRIPT_URL="http://116.205.97.109/scripts/gpu_metrics.sh"  # 脚本下载地址
TARGET_DIR="/opt/scripts"                     # 脚本存放目录
TARGET_SCRIPT="${TARGET_DIR}/gpu_monitor.sh" # 目标脚本路径
CRON_SCHEDULE="*/2 * * * *"                  # 定时任务执行频率
LOG_FILE="/var/log/gpu_monitor.log"          # 日志文件路径
METRICS_DIR="/var/lib/node_exporter/textfile_collector"  # 指标文件目录
METRICS_FILE="${METRICS_DIR}/gpu_metrics.prom"  # 指标文件
NODE_EXPORTER_SERVICE="node_exporter.service"  # node_exporter服务名


# 帮助信息
show_help() {
    echo "用法: $0 [选项]"
    echo "GPU监控脚本部署工具（必须指定--install或--uninstall）"
    echo ""
    echo "选项:"
    echo "  --install    安装并配置GPU监控脚本，重启node_exporter"
    echo "  --uninstall  卸载GPU监控脚本及配置，删除指标文件并重启node_exporter"
    echo "  --help       显示此帮助信息"
    echo ""
}

# 检查是否以root权限运行
check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        echo "错误: 此脚本需要root权限运行，请使用sudo" >&2
        exit 1
    fi
}

# 检查服务是否存在
service_exists() {
    systemctl list-units --all --full -t service --no-legend "$1" | grep -q "$1"
}

# 重启node_exporter并检查状态
restart_node_exporter() {
    echo "检查node_exporter服务状态..."
    if service_exists "${NODE_EXPORTER_SERVICE}"; then
        echo "正在重启node_exporter服务..."
        systemctl restart "${NODE_EXPORTER_SERVICE}" || true  # 允许失败，不中断流程
        
        echo "检查node_exporter状态..."
        if systemctl is-active --quiet "${NODE_EXPORTER_SERVICE}"; then
            echo "node_exporter状态: 运行中"
        else
            echo "警告: node_exporter未正常运行，请查看日志：journalctl -u ${NODE_EXPORTER_SERVICE}" >&2
        fi
    else
        echo "警告: node_exporter服务不存在，跳过重启" >&2
    fi
}

# 安全添加定时任务（避免重复）
add_cron_job() {
    local cron_line="${CRON_SCHEDULE} ${TARGET_SCRIPT} >> ${LOG_FILE} 2>&1"
    
    # 检查任务是否已存在（精确匹配）
    if crontab -l 2>/dev/null | grep -qxF "${cron_line}"; then
        echo "定时任务已存在，跳过添加"
        return 0
    fi
    
    # 备份现有crontab
    local backup_file="/tmp/crontab_backup_$(date +%s)"
    crontab -l 2>/dev/null > "${backup_file}" || true
    
    # 添加新任务
    echo "添加定时任务: ${cron_line}"
    (
        crontab -l 2>/dev/null || true  # 获取现有crontab（如果有）
        echo "${cron_line}"             # 添加新任务
    ) | crontab -
    
    # 验证添加结果
    if crontab -l 2>/dev/null | grep -qxF "${cron_line}"; then
        echo "定时任务添加成功"
    else
        echo "错误: 定时任务添加失败" >&2
        # 恢复备份
        if [ -f "${backup_file}" ]; then
            crontab "${backup_file}"
            echo "已恢复原crontab"
        fi
        return 1
    fi
}

# 安装函数
install_script() {
    echo "===== 开始安装GPU监控脚本 ====="
    
    # 1. 清理旧指标文件（避免残留数据）
    echo "清理旧指标文件..."
    if [ -f "${METRICS_FILE}" ]; then
        rm -f "${METRICS_FILE}"
        echo "已删除旧指标文件: ${METRICS_FILE}"
    fi
    
    # 2. 创建目标目录并下载脚本
    mkdir -p "${TARGET_DIR}"
    cd "${TARGET_DIR}"
    
    echo "正在下载GPU监控脚本..."
    wget -q -O "${TARGET_SCRIPT}" "${SCRIPT_URL}" || {
        echo "错误: 无法从 ${SCRIPT_URL} 下载脚本" >&2
        exit 1
    }
    chmod +x "${TARGET_SCRIPT}"
    
    # 3. 脚本配置（确保GPU总数为8）
    echo "配置脚本参数..."
    sed -i 's/EXPECTED_GPUS=".*"/EXPECTED_GPUS="0 1 2 3 4 5 6 7"/' "${TARGET_SCRIPT}"
    
    # 4. 准备metrics目录
    mkdir -p "${METRICS_DIR}"
    chmod 755 "${METRICS_DIR}"
    
    # 5. 配置定时任务（安全添加，避免重复）
    echo "配置定时任务..."
    add_cron_job || exit 1
    
    # 6. 首次执行脚本生成指标
    echo "测试执行脚本..."
    "${TARGET_SCRIPT}" || {
        echo "错误: 脚本执行失败，请检查日志: ${LOG_FILE}" >&2
        exit 1
    }
    
    # 7. 重启node_exporter使配置生效
    restart_node_exporter
    
    # 8. 验证结果
    echo "===== 安装完成 ====="
    echo "脚本位置: ${TARGET_SCRIPT}"
    echo "定时任务: $(crontab -l | grep "${TARGET_SCRIPT}")"
    echo "日志文件: ${LOG_FILE}"
    echo "指标文件: ${METRICS_FILE}（已生成）"
    echo "node_exporter状态: $(systemctl is-active "${NODE_EXPORTER_SERVICE}" 2>/dev/null || echo "未安装")"
}

# 卸载函数
uninstall_script() {
    echo "===== 开始卸载GPU监控脚本 ====="
    
    # 1. 移除定时任务
    echo "删除定时任务..."
    local cron_line="${CRON_SCHEDULE} ${TARGET_SCRIPT} >> ${LOG_FILE} 2>&1"
    
    if crontab -l 2>/dev/null | grep -qxF "${cron_line}"; then
        # 备份现有crontab
        local backup_file="/tmp/crontab_backup_$(date +%s)"
        crontab -l > "${backup_file}"
        
        # 移除任务
        crontab -l | grep -vxF "${cron_line}" | crontab -
        
        # 验证移除结果
        if ! crontab -l 2>/dev/null | grep -qxF "${cron_line}"; then
            echo "已成功删除定时任务"
        else
            echo "错误: 定时任务删除失败" >&2
            # 恢复备份
            crontab "${backup_file}"
            echo "已恢复原crontab"
            exit 1
        fi
    else
        echo "未找到相关定时任务，跳过"
    fi
    
    # 2. 删除脚本文件
    echo "删除监控脚本..."
    if [ -f "${TARGET_SCRIPT}" ]; then
        rm -f "${TARGET_SCRIPT}"
        echo "已删除脚本: ${TARGET_SCRIPT}"
    else
        echo "脚本文件不存在，跳过"
    fi
    
    # 3. 删除指标文件
    echo "删除指标文件..."
    if [ -f "${METRICS_FILE}" ]; then
        rm -f "${METRICS_FILE}"
        echo "已删除指标文件: ${METRICS_FILE}"
    else
        echo "指标文件不存在，跳过"
    fi
    
    # 4. 删除日志文件
    echo "删除日志文件..."
    if [ -f "${LOG_FILE}" ]; then
        rm -f "${LOG_FILE}"
        echo "已删除日志文件: ${LOG_FILE}"
    else
        echo "日志文件不存在，跳过"
    fi
    
    # 5. 重启node_exporter
    restart_node_exporter
    
    # 6. 验证结果
    echo "===== 卸载完成 ====="
    echo "node_exporter状态: $(systemctl is-active "${NODE_EXPORTER_SERVICE}" 2>/dev/null || echo "未安装")"
}

# 主逻辑
check_root

# 参数检查
if [ $# -ne 1 ]; then
    echo "错误: 必须指定--install或--uninstall" >&2
    show_help
    exit 1
fi

# 执行对应操作
case "$1" in
    --install)
        install_script
        ;;
    --uninstall)
        uninstall_script
        ;;
    --help)
        show_help
        exit 0
        ;;
    *)
        echo "错误: 未知参数 '$1'" >&2
        show_help
        exit 1
        ;;
esac