#!/bin/bash
# ============================================================================
# Sun HPC 系统初始化脚本
# 功能：系统安装完成后自动配置 HPC 计算节点环境
# 版本：1.0
# ============================================================================

set -e

# ============================================================================
# 配置变量（根据实际环境修改）
# ============================================================================

# 日志文件
LOG_FILE="/var/log/sunhpc-init.log"

FRONTEND_IP='@osipaddr@'

# HTTP 服务器地址
HTTP_SERVER="http://${FRONTEND_IP}"

# YUM 仓库配置
YUM_REPO_NAME="sunhpc-local"
YUM_REPO_BASEURL="${HTTP_SERVER}/@ospath@"
YUM_REPO_GPGCHECK=0

# 主机名映射文件（在 HTTP 服务器上）
HOSTNAME_MAP_URL="${HTTP_SERVER}/ks/hostname-map.txt"

# MOTD 文件 URL
MOTD_URL="${HTTP_SERVER}/ks/motd"

# NTP 服务器
NTP_SERVER="ntp.aliyun.com"

# DNS 服务器
DNS_SERVERS="114.114.114.114 223.5.5.5"

# ============================================================================
# 颜色输出函数
# ============================================================================

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 "${GREEN}[INFO]${NC} $(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a $LOG_FILE
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a $LOG_FILE
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a $LOG_FILE
}

log_step() {
    echo -e "${BLUE}[STEP]${NC} $(date '+%Y-%m-%d %H:%M:%S') - $1" | tee -a $LOG_FILE
}

# ============================================================================
# 基础环境准备
# ============================================================================

init_log() {
    # 创建日志文件
    touch $LOG_FILE
    chmod 644 $LOG_FILE
    log_info "=========================================="
    log_info "Sun HPC Node Initialization Started"
    log_info "=========================================="
}

check_network() {
    log_step "Checking network connectivity..."
    
    # 测试网络连接
    if ping -c 1 -W 3 $FRONTEND_IP &>/dev/null; then
        log_info "Network connectivity: OK"
        return 0
    else
        log_warn "Cannot reach HTTP server, will retry later"
        return 1
    fi
}

# ============================================================================
# 1. 创建 .sunhpc-release 文件
# ============================================================================

create_release_file() {
    log_step "Creating .sunhpc-release file..."
    
    cat > /.sunhpc-release << EOF
Sun HPC Platform Release 1.0
Build Date: $(date '+%Y-%m-%d %H:%M:%S')
Hostname: $(hostname)
Kernel: $(uname -r)
Architecture: $(uname -m)
EOF
    
    chmod 644 /.sunhpc-release
    log_info "Created /.sunhpc-release file"
}

# ============================================================================
# 2. 配置主机名和 /etc/hosts
# ============================================================================

configure_hostname() {
    log_step "Configuring hostname and /etc/hosts..."
    
    local mac=$(ip link show | grep -oP 'ether \K[0-9a-f:]+' | head -1 | tr '[:upper:]' '[:lower:]')
    local hostname=""
    
    # 从 HTTP 服务器获取主机名映射
    if check_network; then
        # 下载主机名映射文件
        local map_file="/tmp/hostname-map.txt"
        curl -s -o "$map_file" "$HOSTNAME_MAP_URL" 2>/dev/null
        
        if [[ -f "$map_file" ]]; then
            # 根据 MAC 地址查找主机名
            hostname=$(grep -i "$mac" "$map_file" | awk '{print $2}' | head -1)
            
            # 如果没有找到，根据 IP 最后一段生成
            if [[ -z "$hostname" ]]; then
                local ip_addr=$(ip -4 addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v '127.0.0.1' | head -1)
                local ip_last=$(echo $ip_addr | awk -F'.' '{print $4}')
                hostname="cn$(printf "%03d" $ip_last)"
                log_warn "No hostname mapping for MAC $mac, using generated: $hostname"
            fi
        else
            log_warn "Cannot download hostname map, using default"
            local ip_addr=$(ip -4 addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v '127.0.0.1' | head -1)
            local ip_last=$(echo $ip_addr | awk -F'.' '{print $4}')
            hostname="cn$(printf "%03d" $ip_last)"
        fi
    else
        # 网络不通时使用默认主机名
        hostname="node-$(hostname -s)"
    fi
    
    # 设置主机名
    hostnamectl set-hostname "$hostname"
    log_info "Hostname set to: $hostname"
    
    # 备份原始 hosts 文件
    cp /etc/hosts /etc/hosts.bak.$(date +%Y%m%d)
    
    # 构建新的 hosts 文件
    cat > /etc/hosts << EOF
127.0.0.1   localhost localhost.localdomain localhost4 localhost4.localdomain4
::1         localhost localhost.localdomain localhost6 localhost6.localdomain6

# Sun HPC Node Configuration
$ip_addr   $hostname

# Management Node
$FRONTEND_IP   cluster

# Optional: Add other compute nodes here
EOF
    
    # 尝试下载完整的集群 hosts 文件
    if check_network; then
        local cluster_hosts="${HTTP_SERVER}/ks/cluster-hosts.txt"
        if wget --spider -q "$cluster_hosts" 2>/dev/null; then
            wget -q -O /tmp/cluster-hosts.txt "$cluster_hosts"
            cat /tmp/cluster-hosts.txt >> /etc/hosts
            log_info "Downloaded cluster hosts configuration"
        else
            log_warn "$cluster_hosts not found on server."
        fi
    fi
    
    log_info "Updated /etc/hosts"
}

# ============================================================================
# 3. 配置 MOTD (Message of The Day)
# ============================================================================

configure_motd() {
    log_step "Configuring MOTD..."
    
    # 备份原始 MOTD
    if [[ -f /etc/motd ]]; then
        cp /etc/motd /etc/motd.bak.$(date +%Y%m%d)
    fi
    
    # 尝试从服务器下载 MOTD
    if check_network; then
        if wget --spider -q "$MOTD_URL" 2>/dev/null; then
            wget -q -O /etc/motd "$MOTD_URL"
            log_info "Downloaded MOTD from server"
        else
            log_warn "Cannot download MOTD, creating default"
            create_default_motd
        fi
    else
        create_default_motd
    fi
    
    chmod 644 /etc/motd
}

create_default_motd() {
    cat > /etc/motd << EOF
===========================================================
    Sun HPC Platform - Compute Node
===========================================================
  Welcome to Sun High Performance Computing Cluster

  System Information:
  --------------------
  Hostname....: $(hostname)
  IP Address..: $(ip -4 addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v '127.0.0.1' | head -1)
  OS Version..: $(cat /etc/redhat-release 2>/dev/null || echo "Rocky Linux")
  Kernel......: $(uname -r)
  Architecture: $(uname -m)
  CPU Cores...: $(nproc)
  Memory......: $(free -h | awk '/^Mem:/{print $2}')

  Important URLs:
  --------------------
  • Documentation    : http://172.16.9.254/docs
  • Job Submission   : http://172.16.9.254/slurm
  • Monitoring: http : http://172.16.9.254/monitor

  Slurm Commands:
  --------------------
  sinfo        - View partition information
  squeue       - View job queue
  srun         - Run interactive job
  sbatch       - Submit batch job
  scancel      - Cancel job

===========================================================
  * Unauthorized access is prohibited
  * All activities are logged and monitored
  * For support: hpc@sun.com or ext. 12345
===========================================================
EOF
}

# ============================================================================
# 4. 配置 YUM 本地仓库
# ============================================================================

configure_yum_repo() {
    log_step "Configuring YUM repository..."
    
    local repo_file="/etc/yum.repos.d/${YUM_REPO_NAME}.repo"
    
    # 备份现有 repo 文件
    if [[ -f "$repo_file" ]]; then
        cp "$repo_file" "${repo_file}.bak.$(date +%Y%m%d)"
    fi
    
    # 创建仓库配置文件
    cat > "$repo_file" << EOF
[${YUM_REPO_NAME}]
name=Sun HPC Local Repository
baseurl=${YUM_REPO_BASEURL}
enabled=1
gpgcheck=${YUM_REPO_GPGCHECK}
priority=1
EOF

    # 如果是 Rocky 8/9，可能需要配置 appstream
    if [[ -d "/etc/yum.repos.d" ]]; then
        local appstream_repo="/etc/yum.repos.d/${YUM_REPO_NAME}-appstream.repo"
        cat > "$appstream_repo" << EOF
[${YUM_REPO_NAME}-appstream]
name=Sun HPC Local AppStream Repository
baseurl=${YUM_REPO_BASEURL}-appstream
enabled=1
gpgcheck=${YUM_REPO_GPGCHECK}
priority=1
EOF
    fi
    
    # 清理 YUM 缓存并测试
    yum clean all &>/dev/null
    yum makecache &>/dev/null
    
    if [[ $? -eq 0 ]]; then
        log_info "YUM repository configured successfully"
    else
        log_warn "YUM repository configured but cache generation failed (network may be down)"
    fi
    
    # 显示启用的仓库
    log_info "Enabled repositories:"
    yum repolist 2>/dev/null | grep -E "^${YUM_REPO_NAME}" | tee -a $LOG_FILE
}

# ============================================================================
# 5. 配置 DNS 解析
# ============================================================================

configure_dns() {
    log_step "Configuring DNS resolution..."
    
    # 备份 resolv.conf
    cp /etc/resolv.conf /etc/resolv.conf.bak.$(date +%Y%m%d) 2>/dev/null || true
    
    # 配置 DNS
    cat > /etc/resolv.conf << EOF
# Sun HPC DNS Configuration
nameserver ${DNS_SERVERS%% *}
$(for dns in ${DNS_SERVERS}; do echo "nameserver $dns"; done | tail -n +2)

# Local domain
search sunhpc.local local
EOF
    
    # 防止 NetworkManager 覆盖 resolv.conf
    if [[ -f /etc/NetworkManager/NetworkManager.conf ]]; then
        if ! grep -q "dns=none" /etc/NetworkManager/NetworkManager.conf; then
            sed -i '/^\[main\]/a dns=none' /etc/NetworkManager/NetworkManager.conf
            systemctl restart NetworkManager 2>/dev/null || true
        fi
    fi
    
    log_info "DNS configured: ${DNS_SERVERS}"
}

# ============================================================================
# 6. 配置 NTP 时间同步
# ============================================================================

configure_ntp() {
    log_step "Configuring NTP time synchronization..."
    
    # 使用 chrony (Rocky 8/9 默认)
    if command -v chronyd &>/dev/null; then
        # 备份配置
        cp /etc/chrony.conf /etc/chrony.conf.bak.$(date +%Y%m%d)
        
        # 配置 chrony
        cat > /etc/chrony.conf << EOF
# Sun HPC NTP Configuration
server ${NTP_SERVER} iburst
pool 2.rocky.pool.ntp.org iburst

# Record the rate at which the system clock gains/losses time.
driftfile /var/lib/chrony/drift

# Allow the system clock to be stepped in the first three updates
makestep 1.0 3

# Enable kernel synchronization of the real-time clock (RTC)
rtcsync

# Specify file containing keys for NTP authentication
keyfile /etc/chrony.keys

# Specify directory for log files
logdir /var/log/chrony
EOF
        
        # 重启 chrony
        systemctl restart chronyd &>/dev/null
        systemctl enable chronyd &>/dev/null
        
        log_info "Chrony NTP configured with server: ${NTP_SERVER}"
    else
        # 使用 ntpd (旧系统)
        if command -v ntpd &>/dev/null; then
            cat > /etc/ntp.conf << EOF
server ${NTP_SERVER} iburst
restrict default nomodify notrap nopeer noquery
restrict 127.0.0.1
restrict ::1
driftfile /var/lib/ntp/drift
EOF
            systemctl restart ntpd &>/dev/null
            systemctl enable ntpd &>/dev/null
            log_info "NTP configured with server: ${NTP_SERVER}"
        else
            log_warn "No NTP service found (chronyd or ntpd)"
        fi
    fi
}

# ============================================================================
# 7. 配置 SSH 服务
# ============================================================================

configure_ssh() {
    log_step "Configuring SSH service..."
    
    # 备份配置
    cp /etc/ssh/sshd_config /etc/ssh/sshd_config.bak.$(date +%Y%m%d)
    
    # 优化 SSH 配置
    sed -i 's/#PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config
    sed -i 's/#PubkeyAuthentication.*/PubkeyAuthentication yes/' /etc/ssh/sshd_config
    sed -i 's/#PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config
    sed -i 's/#UseDNS.*/UseDNS no/' /etc/ssh/sshd_config
    sed -i 's/#GSSAPIAuthentication.*/GSSAPIAuthentication no/' /etc/ssh/sshd_config
    
    # 重启 SSH 服务
    systemctl restart sshd &>/dev/null
    
    log_info "SSH service configured"
}

# ============================================================================
# 8. 配置防火墙
# ============================================================================

configure_firewall() {
    log_step "Configuring firewall..."
    
    # 检查 firewalld 是否运行
    if systemctl is-active firewalld &>/dev/null; then
        # 开放常用端口
        firewall-cmd --permanent --add-service=ssh &>/dev/null
        firewall-cmd --permanent --add-service=dhcp &>/dev/null
        firewall-cmd --permanent --add-port=873/tcp &>/dev/null  # rsync
        firewall-cmd --permanent --add-port=111/udp &>/dev/null  # rpcbind
        
        # 重新加载
        firewall-cmd --reload &>/dev/null
        
        log_info "Firewall configured"
    else
        log_warn "Firewalld not running, skipping configuration"
    fi
}

# ============================================================================
# 9. 安装常用软件包
# ============================================================================

install_packages() {
    log_step "Installing common packages..."
    
    local packages=(
        wget curl vim nano
        htop iotop iftop
        net-tools bind-utils
        telnet nc tcpdump
        rsync tree lsof
        gcc make autoconf automake
        openssl-devel zlib-devel
        nfs-utils cifs-utils
        ntpdate
    )
    
    if check_network; then
        for pkg in "${packages[@]}"; do
            if rpm -q "$pkg" &>/dev/null; then
                log_info "Package already installed: $pkg"
            else
                log_info "Installing package: $pkg"
                yum install -y "$pkg" &>> $LOG_FILE || log_warn "Failed to install: $pkg"
            fi
        done
    else
        log_warn "Network unavailable, skipping package installation"
    fi
}

# ============================================================================
# 11. 配置用户环境
# ============================================================================

configure_user_env() {
    log_step "Configuring user environment..."
    
    # 创建共享目录
    mkdir -p /share/{apps,home}
    mkdir -p /state/partition1/{tmp,work}
    
    # 设置权限
    chmod 755 /share
    chmod 1777 /state/partition1/tmp
    
    # 配置全局环境变量
    cat > /etc/profile.d/sunhpc.sh << 'EOF'
# Sun HPC Environment Variables
export SHARE_HOME=/share
export STATE_HOME=/state/partition1
EOF
    
    chmod 644 /etc/profile.d/sunhpc.sh
    log_info "User environment configured"
}

# ============================================================================
# 12. 配置系统优化参数
# ============================================================================

configure_sysctl() {
    log_step "Configuring system optimization..."
    
    cat > /etc/sysctl.d/99-sunhpc.conf << EOF
# Sun HPC System Optimization

# Network optimization
net.core.rmem_max = 134217728
net.core.wmem_max = 134217728
net.ipv4.tcp_rmem = 4096 87380 134217728
net.ipv4.tcp_wmem = 4096 65536 134217728
net.core.netdev_max_backlog = 5000
net.ipv4.tcp_max_syn_backlog = 8192
net.ipv4.tcp_sack = 1
net.ipv4.tcp_timestamps = 1

# Memory optimization
vm.swappiness = 10
vm.dirty_ratio = 30
vm.dirty_background_ratio = 5
vm.vfs_cache_pressure = 50

# File system
fs.file-max = 1048576
fs.inotify.max_user_watches = 1048576

# Process scheduler
kernel.sched_autogroup_enabled = 0
kernel.sched_migration_cost_ns = 5000000
EOF
    
    sysctl -p /etc/sysctl.d/99-sunhpc.conf &>/dev/null
    log_info "System optimization configured"
}

# ============================================================================
# 13. 配置系统限制
# ============================================================================

configure_limits() {
    log_step "Configuring system limits..."
    
    cat > /etc/security/limits.d/99-sunhpc.conf << EOF
# Sun HPC System Limits
* soft nofile 1048576
* hard nofile 1048576
* soft nproc 131072
* hard nproc 131072
* soft memlock unlimited
* hard memlock unlimited
root soft nofile 1048576
root hard nofile 1048576
EOF
    
    log_info "System limits configured"
}

# ============================================================================
# 14. 配置定时任务
# ============================================================================

configure_crontab() {
    log_step "Configuring cron jobs..."
    
    cat > /etc/cron.d/sunhpc << EOF
# Sun HPC Cron Jobs

# Sync time every hour
0 * * * * root /usr/sbin/chronyc -a makestep &>/dev/null

# Clean old logs daily
0 2 * * * root find /var/log -name "*.log" -mtime +30 -delete 2>/dev/null

# Report node status every 5 minutes
*/5 * * * * root /usr/local/bin/node-status-report &>/dev/null
EOF
    
    chmod 644 /etc/cron.d/sunhpc
    log_info "Cron jobs configured"
}

# ============================================================================
# 15. 创建节点状态上报脚本
# ============================================================================

create_report_script() {
    log_step "Creating node status report script..."
    
    cat > /usr/local/bin/node-status-report << 'EOF'
#!/bin/bash
# Node status report script for Sun HPC

HTTP_SERVER="http://172.16.9.254"
NODE_NAME=$(hostname)
IP_ADDR=$(ip -4 addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v '127.0.0.1' | head -1)
LOAD_AVG=$(uptime | awk -F'load average:' '{print $2}' | cut -d, -f1)
MEM_TOTAL=$(free -g | awk '/^Mem:/{print $2}')
MEM_USED=$(free -g | awk '/^Mem:/{print $3}')
DISK_USED=$(df -h / | awk 'NR==2{print $5}' | sed 's/%//')
SLURM_STATUS=$(systemctl is-active slurm-slurmd 2>/dev/null || echo "unknown")

# Build status JSON
STATUS_JSON="{\"node\":\"$NODE_NAME\",\"ip\":\"$IP_ADDR\",\"load\":$LOAD_AVG,\"mem_total\":$MEM_TOTAL,\"mem_used\":$MEM_USED,\"disk_used\":$DISK_USED,\"slurm\":\"$SLURM_STATUS\",\"timestamp\":\"$(date -Iseconds)\"}"

# Send to management node
curl -s -X POST -H "Content-Type: application/json" -d "$STATUS_JSON" ${HTTP_SERVER}/api/node-status &>/dev/null
EOF
    
    chmod +x /usr/local/bin/node-status-report
    log_info "Node status report script created"
}

# ============================================================================
# 16. 配置日志轮转
# ============================================================================

configure_logrotate() {
    log_step "Configuring log rotation..."
    
    cat > /etc/logrotate.d/sunhpc << EOF
/var/log/sunhpc-init.log {
    daily
    rotate 30
    compress
    delaycompress
    missingok
    notifempty
    create 644 root root
}
EOF
    
    log_info "Log rotation configured"
}

# ============================================================================
# 17. 清理临时文件
# ============================================================================

cleanup() {
    log_step "Cleaning up temporary files..."
    
    # 删除临时文件
    rm -f /tmp/hostname-map.txt 2>/dev/null
    rm -f /tmp/cluster-hosts.txt 2>/dev/null
    
    # 清除 YUM 缓存
    yum clean all &>/dev/null
    
    log_info "Cleanup completed"
}

# ============================================================================
# 18. 生成初始化完成报告
# ============================================================================

generate_report() {
    log_step "Generating initialization report..."
    
    cat > /root/init-report.txt << EOF
===========================================
Sun HPC Node Initialization Report
===========================================
Init Time: $(date '+%Y-%m-%d %H:%M:%S')
Hostname: $(hostname)
IP Address: $(ip -4 addr show | grep -oP '(?<=inet\s)\d+(\.\d+){3}' | grep -v '127.0.0.1' | head -1)
MAC Address: $(ip link show | grep -oP 'ether \K[0-9a-f:]+' | head -1)

Services Status:
- NetworkManager: $(systemctl is-active NetworkManager)
- chronyd: $(systemctl is-active chronyd 2>/dev/null || echo "stopped")
- sshd: $(systemctl is-active sshd)
- slurm-slurmd: $(systemctl is-active slurm-slurmd 2>/dev/null || echo "stopped")
- firewalld: $(systemctl is-active firewalld)

YUM Repositories:
$(yum repolist)

Files Created:
- /.sunhpc-release
- /etc/hosts (updated)
- /etc/motd (updated)
- /etc/yum.repos.d/sunhpc-local.repo
- /etc/chrony.conf (updated)
- /etc/security/limits.d/99-sunhpc.conf
- /etc/sysctl.d/99-sunhpc.conf

===========================================
EOF
    
    cat /root/init-report.txt >> $LOG_FILE
    log_info "Initialization report saved to /root/init-report.txt"
}

# ============================================================================
# 主函数
# ============================================================================

main() {
    log_info "Starting Sun HPC node initialization..."
    
    # 执行所有初始化步骤
    init_log
    check_network
    create_release_file
    configure_hostname
    configure_motd
    configure_yum_repo
    configure_dns
    configure_ntp
    configure_ssh
    configure_firewall
    install_packages
    configure_slurm
    configure_user_env
    configure_sysctl
    configure_limits
    configure_crontab
    create_report_script
    configure_logrotate
    cleanup
    generate_report
    
    log_info "=========================================="
    log_info "Sun HPC node initialization completed!"
    log_info "=========================================="
    
    # 显示关键信息
    echo ""
    echo -e "${GREEN}=== Initialization Complete ===${NC}"
    echo -e "Hostname: $(hostname)"
    echo -e "Release file: $(cat /.sunhpc-release | head -1)"
    echo -e "Log file: $LOG_FILE"
    echo -e "Report: /root/init-report.txt"
    echo ""
}

# ============================================================================
# 脚本入口
# ============================================================================

# 检查是否需要重启
if [[ "$1" == "--reboot" ]]; then
    main
    log_info "Rebooting in 5 seconds..."
    sleep 5
    reboot
else
    main
    log_info "Initialization completed. You may reboot if needed."
fi