#!/bin/bash
# 环境部署脚本 - 外部机器接入准备
# 创建人: 乌拉诺斯
# 更新时间: 2025-08-07

set -e

LOG_FILE="/root/deployment.log"
CS_FILE="/root/disk_test"
CACHE_PREFIX="cache"
DOCKER_DATA_DIR="/data0"

# 日志函数
log() {
    echo "$(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a "$LOG_FILE"
}

# 错误处理函数
error_exit() {
    log "错误: $1"
    exit 1
}

# 检查 root 权限
check_root() {
    if [ "$EUID" -ne 0 ]; then
        error_exit "请使用 root 权限运行此脚本"
    fi
}

# 创建目录
create_directories() {
    log "创建必要目录..."
    mkdir -p "$CS_FILE" || error_exit "无法创建目录 $CS_FILE"
    mkdir -p "$DOCKER_DATA_DIR" || error_exit "无法创建目录 $DOCKER_DATA_DIR"
}

# 硬盘挂载和格式化
setup_disks() {
    log "开始设置硬盘..."
    
    local nvme_disk="/dev/nvme0n1"
    local partition="${nvme_disk}p1"  # NVMe分区命名方式不同
    local mount_point="/${CACHE_PREFIX}1"
    
    log "处理NVMe磁盘: $nvme_disk -> $mount_point"
    
    # 检查是否已挂载
    if mount | grep -q "$nvme_disk"; then
        log "磁盘已被挂载，正在卸载..."
        umount "${nvme_disk}"* 2>/dev/null || true
    fi
    
    # 检查是否已有分区
    if [ -e "$partition" ]; then
        log "发现现有分区 $partition，正在清除..."
        wipefs -a "$partition" 2>/dev/null || true
    fi
    
    # 清除磁盘签名
    log "清除磁盘签名..."
    wipefs -a "$nvme_disk" 2>/dev/null || true
    dd if=/dev/zero of="$nvme_disk" bs=1M count=10 2>/dev/null || true
    
    # 创建GPT分区表
    log "创建GPT分区表..."
    parted -s "$nvme_disk" mklabel gpt || {
        log "使用parted创建GPT失败，尝试其他方法..."
        echo "label: gpt" | sfdisk "$nvme_disk"
    }
    
    # 创建分区（NVMe磁盘）
    log "创建分区..."
    parted -s "$nvme_disk" mkpart primary ext4 4096s 100% || {
        log "使用parted创建分区失败，尝试sfdisk..."
        echo ", +" | sfdisk -N 1 "$nvme_disk"
    }
    
    # 等待分区就绪
    log "等待分区就绪..."
    sleep 3
    partprobe "$nvme_disk" 2>/dev/null
    
    # 检查NVMe分区（注意NVMe分区命名是p1, p2等）
    local nvme_partition="${nvme_disk}p1"
    if [ ! -b "$nvme_partition" ]; then
        log "分区设备不存在，尝试重新扫描..."
        echo 1 > "/sys/block/nvme0n1/device/rescan" 2>/dev/null
        sleep 2
    fi
    
    # 格式化分区
    if [ -b "$nvme_partition" ]; then
        log "格式化分区 $nvme_partition ..."
        mkfs.ext4 -F -L "$mount_point" "$nvme_partition" || {
            log "格式化失败，尝试修复..."
            # 尝试多次格式化
            for attempt in {1..3}; do
                sleep 1
                if mkfs.ext4 -F -L "$mount_point" "$nvme_partition"; then
                    break
                fi
            done
        }
        log "格式化成功"
    else
        error_exit "无法找到分区设备"
    fi
    
    # 创建挂载目录
    mkdir -p "$mount_point" || error_exit "无法创建目录 $mount_point"
    
    # 添加到fstab
    log "添加到fstab..."
    sed -i "\|$mount_point|d" /etc/fstab
    echo "LABEL=$mount_point   $mount_point ext4    defaults,nofail,relatime,nodelalloc,nobarrier      0       2" >> /etc/fstab
    
    # 挂载
    log "挂载文件系统..."
    mount "$mount_point" || error_exit "挂载失败"
    
    log "NVMe磁盘设置完成"

}

# 硬盘性能测试
disk_performance_test() {
    log "开始硬盘性能测试..."
    
    # 获取所有cache挂载点
    lsblk | grep "$CACHE_PREFIX" | awk '{print $7}' | sed 's/\///g' > "${CS_FILE}/mycache"
    
    # 运行fio测试
    while read -r line; do
        if [ -n "$line" ]; then
            log "测试硬盘 $line"
            fio -directory="/$line" -direct=1 -iodepth=32 -ioengine=libaio -rw=randrw \
                -rwmixwrite=50 -bs=128k -size=500M -numjobs=10 -nrfiles=10 \
                -runtime=30 -group_reporting -name=fio_test > "${CS_FILE}/fio_${line}" 2>&1
        fi
    done < "${CS_FILE}/mycache"
    
    # 提取IOPS信息
    while read -r line; do
        if [ -n "$line" ]; then
            cat "${CS_FILE}/fio_${line}" | grep read | grep IOPS | awk -F '=|,' '{print $2}' | sed 's/[ \t]*//g' > "${CS_FILE}/IOPS_${line}"
            
            # 检查IOPS是否达标
            IOPS_VALUE=$(cat "${CS_FILE}/IOPS_${line}")
            if [ "$IOPS_VALUE" -lt 1500 ]; then
                log "警告: $line 的IOPS值 $IOPS_VALUE 低于1500"
            else
                log "通过: $line 的IOPS值 $IOPS_VALUE"
            fi
        fi
    done < "${CS_FILE}/mycache"
    
    log "硬盘性能测试完成"
}

# 设置Docker数据目录
setup_docker() {
    log "设置Docker数据目录..."
    
    # 获取数据盘的UUID
    DATA_DISK_UUID=$(blkid -s UUID -o value $(findmnt -n -o SOURCE $DOCKER_DATA_DIR 2>/dev/null) || echo "")
    
    if [ -z "$DATA_DISK_UUID" ]; then
        # 如果没有挂载，使用第一个非系统盘
        DATA_DISK=$(lsblk -d -o NAME,TYPE | grep disk | grep -v "sda\|vda" | head -1 | awk '{print $1}')
        if [ -n "$DATA_DISK" ]; then
            DATA_DISK="/dev/${DATA_DISK}"
            parted -s "$DATA_DISK" mklabel gpt
            parted -s "$DATA_DISK" mkpart primary xfs 4096s 100%
            mkfs.xfs -f "${DATA_DISK}1"
            DATA_DISK_UUID=$(blkid -s UUID -o value "${DATA_DISK}1")
            echo "UUID=${DATA_DISK_UUID}   $DOCKER_DATA_DIR  xfs  defaults  0  0" >> /etc/fstab
            mount -a
        fi
    fi
    
    # 配置Docker
    cat > /etc/docker/daemon.json << EOF
{
  "data-root": "$DOCKER_DATA_DIR"
}
EOF
    
    log "Docker配置完成"
}

# 检查端口占用
check_ports() {
    log "检查端口占用..."
    PORTS=(57891 63501 65501)
    
    for port in "${PORTS[@]}"; do
        if netstat -tuln | grep ":$port " > /dev/null; then
            error_exit "端口 $port 已被占用"
        fi
    done
    log "端口检查通过"
}

# 主执行函数
main() {
    log "开始环境部署..."
    check_root
    create_directories
    check_ports
    setup_disks
    disk_performance_test
    setup_docker
    
    log "环境部署完成！请检查以上输出是否有错误。"
    log "接下来请运行业务部署脚本。"
}

# 执行主函数
main "$@"