#!/bin/bash
# ServerGuard 安装脚本
# 支持 Debian/Ubuntu 和 CentOS/RHEL

echo "========================================"
echo "ServerGuard 安装脚本"
echo "========================================"
echo ""

# 检查是否为 root
if [ "$EUID" -ne 0 ]; then 
    echo "错误: 请以 root 权限运行此脚本"
    echo "  sudo ./install.sh"
    exit 1
fi

# 检测 Linux 发行版
if [ -f /etc/os-release ]; then
    . /etc/os-release
    OS=$NAME
    VER=$VERSION_ID
else
    echo "无法检测操作系统类型"
    exit 1
fi

echo "检测到操作系统: $OS $VER"
echo ""

# 记录安装失败的包
FAILED_PACKAGES=""

# 安装单个包的函数
install_package() {
    local pkg=$1
    local pkg_manager=$2
    
    if [ "$pkg_manager" = "apt" ]; then
        apt-get install -y "$pkg" 2>/dev/null && return 0
    else
        yum install -y "$pkg" 2>/dev/null && return 0
    fi
    
    FAILED_PACKAGES="$FAILED_PACKAGES $pkg"
    return 1
}

# 安装 Debian/Ubuntu 依赖
install_debian_deps() {
    echo "正在安装 Debian/Ubuntu 依赖..."
    apt-get update
    
    # 核心依赖（必须）
    CORE_PKGS="lshw dmidecode smartmontools lm-sensors ipmitool mdadm pciutils usbutils util-linux coreutils grep gawk sed"
    
    # 可选依赖
    OPTIONAL_PKGS="stress-ng memtester edac-utils fio nvme-cli"
    
    echo "安装核心依赖..."
    for pkg in $CORE_PKGS; do
        install_package "$pkg" "apt" || echo "警告: $pkg 安装失败"
    done
    
    echo "安装可选依赖..."
    for pkg in $OPTIONAL_PKGS; do
        install_package "$pkg" "apt" || echo "注意: $pkg 安装失败（可选）"
    done
}

# 安装 RHEL/CentOS 依赖
install_redhat_deps() {
    echo "正在安装 RHEL/CentOS 依赖..."
    
    # 尝试启用 EPEL
    if ! rpm -qa | grep -q epel-release; then
        echo "启用 EPEL 仓库..."
        yum install -y epel-release 2>/dev/null || true
    fi
    
    # 对于 CentOS 8/RHEL 8，启用 PowerTools/CRB 仓库
    if [[ "$VER" == 8* ]] || [[ "$VER" == "8" ]]; then
        echo "启用 PowerTools 仓库..."
        yum config-manager --set-enabled powertools 2>/dev/null || \
        yum config-manager --set-enabled PowerTools 2>/dev/null || true
        
        # 尝试启用 CRB (CodeReady Builder) 对于 RHEL 8
        subscription-manager repos --enable codeready-builder-for-rhel-8-x86_64-rpms 2>/dev/null || true
    fi
    
    # 核心依赖（必须）
    CORE_PKGS="lshw dmidecode smartmontools lm_sensors ipmitool mdadm pciutils usbutils util-linux coreutils grep gawk sed"
    
    echo "安装核心依赖..."
    for pkg in $CORE_PKGS; do
        install_package "$pkg" "yum" || echo "警告: $pkg 安装失败"
    done
    
    # 尝试安装 OpenIPMI (替代 ipmitool 的依赖)
    install_package "OpenIPMI" "yum" || echo "注意: OpenIPMI 安装失败（可选）"
    
    # 可选依赖
    OPTIONAL_PKGS="memtester edac-utils fio nvme-cli"
    
    echo "安装可选依赖..."
    for pkg in $OPTIONAL_PKGS; do
        install_package "$pkg" "yum" || echo "注意: $pkg 安装失败（可选）"
    done
    
    # 特别处理 stress-ng
    echo "尝试安装 stress-ng..."
    if ! yum install -y stress-ng 2>/dev/null; then
        echo "注意: stress-ng 从默认仓库安装失败"
        
        # 尝试从 EPEL 安装 stress (备选)
        echo "尝试安装 stress 作为备选..."
        if yum install -y stress 2>/dev/null; then
            echo "stress 安装成功，可作为压力测试备选工具"
        else
            echo "警告: stress 和 stress-ng 都安装失败"
            echo "      压力测试功能将不可用"
            FAILED_PACKAGES="$FAILED_PACKAGES stress-ng"
        fi
    fi
    
    # 对于 CentOS 8，提供手动安装 stress-ng 的指导
    if [[ "$VER" == 8* ]] && [[ "$FAILED_PACKAGES" == *"stress-ng"* ]]; then
        echo ""
        echo "============================================"
        echo "注意: CentOS 8 中 stress-ng 需要从源码编译安装"
        echo "============================================"
        echo "手动安装步骤:"
        echo "  1. 安装编译依赖:"
        echo "     yum install -y gcc make libaio-devel libattr-devel libbsd-devel libcap-devel libgcrypt-devel"
        echo "  2. 下载并编译 stress-ng:"
        echo "     cd /tmp"
        echo "     git clone https://github.com/ColinIanKing/stress-ng.git"
        echo "     cd stress-ng"
        echo "     make"
        echo "     make install"
        echo "============================================"
        echo ""
    fi
}

# 根据发行版安装
case "$OS" in
    *Debian*|*Ubuntu*)
        install_debian_deps
        ;;
    *CentOS*|*Red*Hat*|*Fedora*|*Alma*|*Rocky*)
        install_redhat_deps
        ;;
    *)
        echo "不支持的操作系统: $OS"
        echo "请手动安装以下工具:"
        echo "  lshw, dmidecode, smartmontools, lm-sensors, stress-ng, memtester"
        echo "  ipmitool, edac-utils, fio, mdadm, pciutils, usbutils"
        exit 1
        ;;
esac

echo ""
echo "系统依赖安装完成"

# 显示安装失败的包
if [ -n "$FAILED_PACKAGES" ]; then
    echo ""
    echo "以下包安装失败: $FAILED_PACKAGES"
    echo "某些功能可能受限，ServerGuard 仍可运行基本检测"
fi

echo ""

# 检查 Python 版本
echo "检查 Python 版本..."
if command -v python3 &> /dev/null; then
    PYTHON_VERSION=$(python3 --version 2>&1 | awk '{print $2}')
    echo "找到 Python $PYTHON_VERSION"
elif command -v python &> /dev/null; then
    PYTHON_VERSION=$(python --version 2>&1 | awk '{print $2}')
    echo "找到 Python $PYTHON_VERSION"
else
    echo "错误: 未找到 Python"
    echo "请安装 Python 3.6 或更高版本"
    exit 1
fi

# 检查 Python 版本号
PYTHON_MAJOR=$(echo $PYTHON_VERSION | cut -d. -f1)
PYTHON_MINOR=$(echo $PYTHON_VERSION | cut -d. -f2)

if [ "$PYTHON_MAJOR" -lt 3 ] || ([ "$PYTHON_MAJOR" -eq 3 ] && [ "$PYTHON_MINOR" -lt 6 ]); then
    echo "错误: Python 版本过低 ($PYTHON_VERSION)"
    echo "需要 Python 3.6 或更高版本"
    exit 1
fi

echo "Python 版本符合要求"
echo ""

# 安装 Python 依赖
echo "安装 Python 依赖..."
PIP_CMD="pip3"
if ! command -v pip3 &> /dev/null; then
    PIP_CMD="pip"
fi

$PIP_CMD install -r requirements.txt || {
    echo "警告: pip 安装失败，尝试使用 --user 选项"
    $PIP_CMD install --user -r requirements.txt
}

echo ""

# 配置 lm-sensors
if command -v sensors-detect &> /dev/null; then
    echo ""
    echo "检测到 lm-sensors 需要配置"
    echo "是否要运行 sensors-detect 配置传感器? (y/N)"
    read -r response
    if [[ "$response" =~ ^([yY][eE][sS]|[yY])$ ]]; then
        echo "正在运行 sensors-detect..."
        sensors-detect --auto || true
    fi
fi

echo ""
echo "========================================"
echo "安装完成!"
echo "========================================"
echo ""

# 检查依赖状态
echo "依赖检查:"
echo "------------"
for cmd in lshw dmidecode smartctl sensors ipmitool; do
    if command -v "$cmd" &> /dev/null; then
        echo "  ✓ $cmd"
    else
        echo "  ✗ $cmd (未安装)"
    fi
done

echo ""
echo "压力测试工具:"
if command -v stress-ng &> /dev/null; then
    echo "  ✓ stress-ng (推荐)"
elif command -v stress &> /dev/null; then
    echo "  ✓ stress (备选)"
else
    echo "  ✗ stress/stress-ng (未安装，压力测试不可用)"
fi

echo ""
echo "使用方法:"
echo "  快速检测:    sudo python3 main.py --quick"
echo "  全面诊断:    sudo python3 main.py --full"
echo "  特定模块:    sudo python3 main.py --module cpu"
echo "  生成报告:    sudo python3 main.py --quick --format json --output report.json"
echo ""
echo "查看帮助: python3 main.py --help"
echo ""

# 创建快捷方式（可选）
echo "是否要创建 /usr/local/bin/serverguard 快捷方式? (y/N)"
read -r response
if [[ "$response" =~ ^([yY][eE][sS]|[yY])$ ]]; then
    SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
    cat > /usr/local/bin/serverguard << EOF
#!/bin/bash
cd "$SCRIPT_DIR"
python3 main.py "\$@"
EOF
    chmod +x /usr/local/bin/serverguard
    echo "快捷方式已创建: serverguard"
    echo "现在可以直接使用: sudo serverguard --quick"
fi

echo ""
echo "安装完成!"