#!/bin/bash

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 更新安装记录
update_install_record() {
    local pid="$1"
    # 使用传入的安装目录参数，如果没有则使用默认值
    local install_base_dir="${2:-/opt/argus-metric/current}"
    local install_record="$install_base_dir/.install_record"
    
    # 如果安装记录文件不存在，说明是首次安装，由主安装脚本统一创建
    if [[ ! -f "$install_record" ]]; then
        log_info "安装记录文件不存在，将由主安装脚本创建"
        return 0
    fi
    
    # 如果文件存在，说明是重启场景，只更新 PID 字段
    if command -v jq &> /dev/null; then
        # 读取当前 PID
        local current_pid=$(jq -r '.components."dcgm-exporter".pid // ""' "$install_record" 2>/dev/null)
        
        if [[ -z "$current_pid" ]]; then
            log_warning "无法读取当前 PID，跳过更新"
            return 1
        fi
        
        # 使用 jq 只更新 pid 字段，保持字符串类型，保留其他字段
        jq --arg new_pid "$pid" '.components."dcgm-exporter".pid = $new_pid' "$install_record" > "$install_record.tmp" && mv "$install_record.tmp" "$install_record"
        log_info "PID 已更新: $current_pid -> $pid"
    else
        log_warning "jq 命令不可用，无法更新安装记录文件"
    fi
}

# 显示帮助信息
show_help() {
    echo "DCGM Exporter 安装脚本"
    echo
    echo "用法: $0 [选项]"
    echo
    echo "选项:"
    echo "  --help     显示此帮助信息"
    echo
    echo "示例:"
    echo "  $0          # 安装 DCGM Exporter"
    echo
}

# 解析命令行参数
INSTALL_DIR=""
for arg in "$@"; do
    case $arg in
        --help|-h)
            show_help
            exit 0
            ;;
        *)
            # 如果参数不是以--开头，则认为是安装目录
            if [[ ! "$arg" =~ ^-- ]]; then
                INSTALL_DIR="$arg"
            else
                log_error "未知参数: $arg"
                show_help
                exit 1
            fi
            ;;
    esac
done

# 检查是否为 root 用户
check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "此脚本需要 root 权限运行"
        log_info "请使用: sudo $0"
        exit 1
    fi
}

# 检查系统要求
check_system() {
    log_info "检查系统要求..."
    
    # 检查操作系统
    if [[ ! -f /etc/os-release ]]; then
        log_error "无法检测操作系统版本"
        exit 1
    fi
    
    source /etc/os-release
    log_info "检测到操作系统: $NAME $VERSION"
    
    # 检查是否为 Ubuntu/Debian
    if [[ "$ID" != "ubuntu" && "$ID" != "debian" ]]; then
        log_warning "此脚本主要针对 Ubuntu/Debian 系统，其他系统可能需要调整"
    fi
    
    # 检查 NVIDIA GPU
    if ! command -v nvidia-smi &> /dev/null; then
        log_warning "未检测到 nvidia-smi，请确保已安装 NVIDIA 驱动"
    else
        log_success "检测到 NVIDIA GPU"
        nvidia-smi --query-gpu=name --format=csv,noheader,nounits | head -1
    fi
}

# 安装 DCGM 依赖
install_dcgm_dependency() {
    log_info "安装 DCGM 依赖..."
    
    local deb_file="bin/datacenter-gpu-manager_3.3.9_amd64.deb"
    
    if [[ ! -f "$deb_file" ]]; then
        log_error "找不到 DCGM 依赖文件: $deb_file"
        exit 1
    fi
    
    # 安装 deb 包
    dpkg -i "$deb_file" || {
        log_warning "dpkg 安装失败，尝试使用 apt 修复依赖..."
        apt-get update
        apt-get install -f -y
        dpkg -i "$deb_file"
    }
    
    log_success "DCGM 依赖安装完成"
}

# 检查 DCGM 服务状态
check_dcgm_service() {
    log_info "检查 DCGM 服务状态..."
    
    # 检查 DCGM 服务是否在运行
    if systemctl is-active --quiet dcgm 2>/dev/null; then
        log_success "DCGM 服务已在运行"
    elif pgrep -f nv-hostengine > /dev/null; then
        log_success "nv-hostengine 进程已在运行"
    else
        log_warning "DCGM 服务未运行，需要手动启动"
        log_info "启动 DCGM 服务的方法:"
        log_info "  1. 使用 systemd: sudo systemctl start dcgm"
        log_info "  2. 手动启动: nohup nv-hostengine > /var/log/nv-hostengine.log 2>&1 &"
    fi
    
    # 测试 DCGM 连接
    if systemctl is-active --quiet dcgm 2>/dev/null || pgrep -f nv-hostengine > /dev/null; then
        log_info "测试 DCGM 连接..."
        if dcgmi discovery -l > /dev/null 2>&1; then
            log_success "DCGM 连接测试成功"
        else
            log_warning "DCGM 连接测试失败，请检查服务状态"
        fi
    fi
}

# 停止可能运行的服务
stop_existing_service() {
    log_info "检查并停止可能运行的服务..."
    
    local pid_file="/var/run/dcgm-exporter.pid"
    
    # 检查并停止通过 PID 文件管理的服务
    if [[ -f "$pid_file" ]]; then
        local pid=$(cat "$pid_file")
        if kill -0 "$pid" 2>/dev/null; then
            log_info "发现正在运行的 DCGM Exporter 服务 (PID: $pid)，正在停止..."
            kill "$pid" > /dev/null 2>&1 || true
            sleep 2
            if kill -0 "$pid" 2>/dev/null; then
                log_warning "进程未响应，强制终止..."
                kill -9 "$pid" > /dev/null 2>&1 || true
            fi
            rm -f "$pid_file"
            log_success "服务已停止"
        else
            log_warning "发现过期的 PID 文件，正在清理..."
            rm -f "$pid_file"
        fi
    fi
    
    # 查找并停止所有 dcgm-exporter 进程（排除脚本自身）
    local exporter_bin="/usr/local/bin/dcgm-exporter"
    local pids=$(pgrep -f "$exporter_bin")

    if [[ -n "$pids" ]]; then
        log_info "发现其他 dcgm-exporter 进程，正在停止..."
        for pid in $pids; do
            if [[ "$pid" != "$$" ]]; then
                kill "$pid" > /dev/null 2>&1 || true
                sleep 1
                if kill -0 "$pid" 2>/dev/null; then
                    log_warning "进程 $pid 未响应，强制终止..."
                    kill -9 "$pid" > /dev/null 2>&1 || true
                fi
            fi
        done
        log_success "所有 dcgm-exporter 进程已停止"
    fi
}

# 安装 DCGM Exporter 二进制文件
install_dcgm_exporter() {
    log_info "安装 DCGM Exporter..."
    
    local binary_file="bin/dcgm-exporter"
    local install_dir="/usr/local/bin"
    
    if [[ ! -f "$binary_file" ]]; then
        log_error "找不到 DCGM Exporter 二进制文件: $binary_file"
        exit 1
    fi
    
    # 停止可能运行的服务
    stop_existing_service
    
    # 复制二进制文件
    cp "$binary_file" "$install_dir/"
    chmod +x "$install_dir/dcgm-exporter"
    
    log_success "DCGM Exporter 二进制文件安装完成"
}

# 安装配置文件
install_config() {
    log_info "安装配置文件..."
    
    local config_dir="/etc/dcgm-exporter"
    local config_file="config/default-counters.csv"
    
    # 创建配置目录
    mkdir -p "$config_dir"
    
    if [[ -f "$config_file" ]]; then
        cp "$config_file" "$config_dir/"
        log_success "配置文件安装完成"
    else
        log_warning "未找到配置文件，使用默认配置"
    fi
}

# 启动 DCGM Exporter 服务
start_dcgm_exporter() {
    log_info "启动 DCGM Exporter 服务..."
    
    local binary_path="/usr/local/bin/dcgm-exporter"
    local log_file="/var/log/dcgm-exporter.log"
    local pid_file="/var/run/dcgm-exporter.pid"
    
    # 检查服务是否已经在运行
    if [[ -f "$pid_file" ]]; then
        local pid=$(cat "$pid_file")
        if kill -0 "$pid" 2>/dev/null; then
            log_info "DCGM Exporter 服务已在运行 (PID: $pid)"
            return 0
        else
            log_warning "发现过期的 PID 文件，正在清理..."
            rm -f "$pid_file"
        fi
    fi
    
    # 检查端口是否被占用
    if netstat -tuln 2>/dev/null | grep -q ":9400 "; then
        log_warning "端口 9400 已被占用，请检查是否有其他服务在运行"
        return 1
    fi
    
    # 启动服务
    log_info "正在启动 DCGM Exporter..."
    nohup "$binary_path" --address=:9400 > "$log_file" 2>&1 &
    local pid=$!
    
    # 保存 PID
    echo "$pid" > "$pid_file"
    
    # 等待服务启动
    sleep 2
    
    # 检查服务是否成功启动
    if kill -0 "$pid" 2>/dev/null; then
        log_success "DCGM Exporter 服务启动成功 (PID: $pid)"
        log_info "日志文件: $log_file"
        log_info "PID 文件: $pid_file"
        
        # 更新安装记录
        update_install_record "$pid" "$INSTALL_DIR"
    else
        log_error "DCGM Exporter 服务启动失败"
        rm -f "$pid_file"
        return 1
    fi
}



# 显示安装信息
show_install_info() {
    log_success "DCGM Exporter 安装完成！"
    echo
    echo "安装信息:"
    echo "  二进制文件: /usr/local/bin/dcgm-exporter"
    echo "  配置文件: /etc/dcgm-exporter/default-counters.csv"
    echo "  默认端口: 9400"
    echo
    echo "使用方法:"
    echo "  1. 启动 DCGM 服务:"
    echo "     sudo systemctl start dcgm"
    echo "     或: nohup nv-hostengine > /var/log/nv-hostengine.log 2>&1 &"
    echo "  2. 启动 DCGM Exporter:"
    echo "     /usr/local/bin/dcgm-exporter --address=:9400"
    echo "     或: nohup /usr/local/bin/dcgm-exporter --address=:9400 &"
    echo
    echo "测试连接:"
    echo "  curl http://localhost:9400/metrics"
    echo
}

# 主函数
main() {
    echo "=========================================="
    echo "    DCGM Exporter 安装脚本 v1.0"
    echo "=========================================="
    echo
    
    check_root
    check_system
    
    log_info "开始安装 DCGM Exporter..."
    
    install_dcgm_dependency
    check_dcgm_service
    install_dcgm_exporter
    install_config
    start_dcgm_exporter
    
    show_install_info
}

# 脚本入口
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi