#!/bin/bash

################################################################################
# Linux 系统初始化脚本 v1.3
# 功能: 镜像源切换 | 基础工具安装 | 系统优化配置
# 支持: Ubuntu/Debian, CentOS/RHEL, Rocky, AlmaLinux
# 作者: 自动化运维工具
# 使用: sudo ./init_system.sh [--auto] [--mirror=aliyun] [--non-interactive]
################################################################################

set -e

# 颜色定义 - 用于终端彩色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
PURPLE='\033[0;35m'
NC='\033[0m'  # No Color

# 全局变量
OS=""
OS_NAME=""
OS_VERSION=""

# 输出函数定义
print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

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

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

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

# 脚本横幅显示
print_banner() {
    clear
    echo -e "${PURPLE}"
    echo "=================================================="
    echo "       Linux 系统初始化脚本 v1.3"
    echo "=================================================="
    echo "功能: 镜像源切换 | 基础工具安装 | 系统优化配置"
    echo "支持: Ubuntu/Debian, CentOS/RHEL, Rocky, AlmaLinux"
    echo "=================================================="
    echo -e "${NC}"
}

# 检查脚本运行权限
check_permissions() {
    if [[ $EUID -ne 0 ]]; then
        print_error "此脚本需要 root 权限"
        print_info "请使用: sudo $0"
        exit 1
    fi
    print_success "权限检查通过"
}

# 检测操作系统类型和版本
detect_os() {
    print_info "正在检测操作系统..."
    
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS="$ID"
        OS_NAME="$NAME"
        OS_VERSION="${VERSION_ID:-unknown}"
        print_success "操作系统: $OS_NAME $OS_VERSION"
    else
        print_error "无法检测操作系统"
        exit 1
    fi
}

# 备份原始软件源
backup_sources() {
    print_info "正在备份原始软件源..."
    
    local timestamp=$(date +%Y%m%d-%H%M%S)
    
    case $OS in
        ubuntu|debian)
            if [[ -f /etc/apt/sources.list ]]; then
                cp /etc/apt/sources.list /etc/apt/sources.list.bak.$timestamp
                print_success "已备份到 /etc/apt/sources.list.bak.$timestamp"
            fi
            ;;
            
        centos|rhel|rocky|almalinux)
            if [[ -d /etc/yum.repos.d ]]; then
                mkdir -p /etc/yum.repos.d/backup-$timestamp
                cp /etc/yum.repos.d/*.repo /etc/yum.repos.d/backup-$timestamp/ 2>/dev/null || true
                print_success "已备份到 /etc/yum.repos.d/backup-$timestamp/"
            fi
            ;;
    esac
}

# 配置 Ubuntu/Debian 镜像源
configure_ubuntu_mirrors() {
    print_info "正在配置 Ubuntu/Debian 镜像源..."
    
    # 获取 Ubuntu 版本代号
    local ubuntu_codename
    if command -v lsb_release &> /dev/null; then
        ubuntu_codename=$(lsb_release -cs)
    else
        # 根据版本号推测代号
        case "$OS_VERSION" in
            "22.04") ubuntu_codename="jammy" ;;
            "20.04") ubuntu_codename="focal" ;;
            "18.04") ubuntu_codename="bionic" ;;
            "16.04") ubuntu_codename="xenial" ;;
            *) ubuntu_codename="jammy" ;;
        esac
    fi
    
    # 如果指定了镜像源参数，直接使用
    local mirror_url=""
    local mirror_name=""
    
    if [[ -n "${MIRROR_ARG:-}" ]]; then
        case "${MIRROR_ARG}" in
            aliyun)
                mirror_url="http://mirrors.aliyun.com"
                mirror_name="阿里云"
                ;;
            tsinghua)
                mirror_url="https://mirrors.tuna.tsinghua.edu.cn"
                mirror_name="清华大学"
                ;;
            ustc)
                mirror_url="https://mirrors.ustc.edu.cn"
                mirror_name="中国科技大学"
                ;;
            *)
                print_warning "未知的镜像源参数，使用交互式选择"
                ;;
        esac
    fi
    
    # 交互式选择镜像源
    if [[ -z "$mirror_url" ]] && [[ -z "${NON_INTERACTIVE:-}" ]]; then
        echo ""
        echo "请选择镜像源:"
        echo "1) 阿里云 (推荐)"
        echo "2) 清华大学"
        echo "3) 中国科技大学"
        echo "4) 官方源"
        echo ""
        
        read -p "请选择 [1-4, 默认: 1]: " mirror_choice
        mirror_choice=${mirror_choice:-1}
        
        case $mirror_choice in
            1)
                mirror_url="http://mirrors.aliyun.com"
                mirror_name="阿里云"
                ;;
            2)
                mirror_url="https://mirrors.tuna.tsinghua.edu.cn"
                mirror_name="清华大学"
                ;;
            3)
                mirror_url="https://mirrors.ustc.edu.cn"
                mirror_name="中国科技大学"
                ;;
            4)
                mirror_url="http://archive.ubuntu.com"
                mirror_name="官方源"
                ;;
            *)
                mirror_url="http://mirrors.aliyun.com"
                mirror_name="阿里云"
                ;;
        esac
    elif [[ -z "$mirror_url" ]]; then
        # 非交互模式默认使用阿里云
        mirror_url="http://mirrors.aliyun.com"
        mirror_name="阿里云"
    fi
    
    print_info "正在配置 $mirror_name 镜像源..."
    
    # 生成新的 sources.list
    if [[ "$mirror_url" == "http://archive.ubuntu.com" ]]; then
        cat > /etc/apt/sources.list <<EOF
# Ubuntu 官方源
deb http://archive.ubuntu.com/ubuntu/ $ubuntu_codename main restricted universe multiverse
deb http://archive.ubuntu.com/ubuntu/ $ubuntu_codename-updates main restricted universe multiverse
deb http://archive.ubuntu.com/ubuntu/ $ubuntu_codename-backports main restricted universe multiverse
deb http://security.ubuntu.com/ubuntu/ $ubuntu_codename-security main restricted universe multiverse
EOF
    else
        cat > /etc/apt/sources.list <<EOF
# $mirror_name 镜像源
deb $mirror_url/ubuntu/ $ubuntu_codename main restricted universe multiverse
deb $mirror_url/ubuntu/ $ubuntu_codename-updates main restricted universe multiverse
deb $mirror_url/ubuntu/ $ubuntu_codename-backports main restricted universe multiverse
deb $mirror_url/ubuntu/ $ubuntu_codename-security main restricted universe multiverse
EOF
    fi
    
    print_success "$mirror_name 镜像源配置完成"
    
    # 更新软件包索引
    print_info "正在更新软件包索引..."
    apt-get update
    print_success "软件包索引更新完成"
}

# 配置 CentOS/RHEL/Rocky/AlmaLinux 镜像源
configure_centos_mirrors() {
    print_info "正在配置 CentOS/RHEL/Rocky/AlmaLinux 镜像源..."
    
    local mirror_url=""
    local mirror_name=""
    
    if [[ -n "${MIRROR_ARG:-}" ]]; then
        case "${MIRROR_ARG}" in
            aliyun)
                mirror_url="https://mirrors.aliyun.com"
                mirror_name="阿里云"
                ;;
            tsinghua)
                mirror_url="https://mirrors.tuna.tsinghua.edu.cn"
                mirror_name="清华大学"
                ;;
            ustc)
                mirror_url="https://mirrors.ustc.edu.cn"
                mirror_name="中国科技大学"
                ;;
            *)
                print_warning "未知的镜像源参数，使用交互式选择"
                ;;
        esac
    fi
    
    # 交互式选择镜像源
    if [[ -z "$mirror_url" ]] && [[ -z "${NON_INTERACTIVE:-}" ]]; then
        echo ""
        echo "请选择镜像源:"
        echo "1) 阿里云 (推荐)"
        echo "2) 清华大学"
        echo "3) 中国科技大学"
        echo "4) 官方源"
        echo ""
        
        read -p "请选择 [1-4, 默认: 1]: " mirror_choice
        mirror_choice=${mirror_choice:-1}
        
        case $mirror_choice in
            1)
                mirror_url="https://mirrors.aliyun.com"
                mirror_name="阿里云"
                ;;
            2)
                mirror_url="https://mirrors.tuna.tsinghua.edu.cn"
                mirror_name="清华大学"
                ;;
            3)
                mirror_url="https://mirrors.ustc.edu.cn"
                mirror_name="中国科技大学"
                ;;
            4)
                print_info "使用官方源"
                mirror_name="官方源"
                ;;
            *)
                mirror_url="https://mirrors.aliyun.com"
                mirror_name="阿里云"
                ;;
        esac
    elif [[ -z "$mirror_url" ]]; then
        mirror_url="https://mirrors.aliyun.com"
        mirror_name="阿里云"
    fi
    
    print_info "正在配置 $mirror_name 镜像源..."
    
    # 根据不同的发行版配置
    if [[ "$mirror_name" != "官方源" ]]; then
        case $OS in
            centos)
                if [[ "$OS_VERSION" == "7" ]]; then
                    sed -i.bak -e "s|^mirrorlist=|#mirrorlist=|g" \
                             -e "s|^#baseurl=http://mirror.centos.org|baseurl=$mirror_url/centos|g" \
                             /etc/yum.repos.d/CentOS-*.repo
                elif [[ "$OS_VERSION" == "8" ]]; then
                    # CentOS 8 已停止支持，使用 vault 源
                    sed -i.bak -e "s|^mirrorlist=|#mirrorlist=|g" \
                             -e "s|^#baseurl=http://mirror.centos.org|baseurl=$mirror_url/centos-vault|g" \
                             /etc/yum.repos.d/CentOS-*.repo
                fi
                ;;
                
            rocky|almalinux)
                if [[ "$OS_VERSION" == "8"* ]]; then
                    sed -i.bak -e "s|^mirrorlist=|#mirrorlist=|g" \
                             -e "s|^#baseurl=http://dl.rockylinux.org|baseurl=$mirror_url/rocky|g" \
                             /etc/yum.repos.d/Rocky-*.repo /etc/yum.repos.d/almalinux-*.repo 2>/dev/null || true
                elif [[ "$OS_VERSION" == "9"* ]]; then
                    sed -i.bak -e "s|^mirrorlist=|#mirrorlist=|g" \
                             -e "s|^#baseurl=http://dl.rockylinux.org|baseurl=$mirror_url/rocky|g" \
                             /etc/yum.repos.d/rocky-*.repo /etc/yum.repos.d/almalinux-*.repo 2>/dev/null || true
                fi
                ;;
        esac
    fi
    
    # 清理缓存并更新
    print_info "清理缓存并更新..."
    yum clean all
    yum makecache
    
    print_success "镜像源配置完成"
}

# 安装基础工具
install_basic_tools() {
    print_info "正在安装基础工具..."
    
    # 工具列表
    local ubuntu_tools=(
        curl wget git
        vim nano
        htop
        net-tools
        unzip zip tar
        build-essential
        software-properties-common
        apt-transport-https
        ca-certificates
        gnupg
        lsb-release
    )
    
    local centos_tools=(
        curl wget git
        vim nano
        htop
        net-tools
        unzip zip tar
        gcc gcc-c++ make
        epel-release
    )
    
    # 询问用户是否安装
    if [[ -z "${NON_INTERACTIVE:-}" ]]; then
        echo ""
        read -p "是否安装基础工具? [Y/n]: " install_tools
        install_tools=${install_tools:-Y}
        
        if [[ ! $install_tools =~ ^[Yy]$ ]]; then
            print_info "跳过基础工具安装"
            return
        fi
    fi
    
    # 安装工具
    case $OS in
        ubuntu|debian)
            apt-get update
            apt-get install -y "${ubuntu_tools[@]}"
            ;;
            
        centos|rhel|rocky|almalinux)
            if [[ "$OS" == "centos" ]] || [[ "$OS" == "rhel" ]]; then
                yum install -y epel-release || true
            fi
            yum install -y "${centos_tools[@]}"
            ;;
    esac
    
    print_success "基础工具安装完成"
}

# 配置 SSH 服务
configure_ssh() {
    print_info "正在配置 SSH 服务..."
    
    # 检查 SSH 是否已安装
    if ! command -v sshd &> /dev/null; then
        print_warning "SSH 服务未安装"
        case $OS in
            ubuntu|debian)
                apt-get install -y openssh-server
                ;;
            centos|rhel|rocky|almalinux)
                yum install -y openssh-server
                ;;
        esac
    fi
    
    if [[ -z "${NON_INTERACTIVE:-}" ]]; then
        read -p "是否优化 SSH 配置? [y/N]: " config_ssh
        config_ssh=${config_ssh:-N}
        
        if [[ ! $config_ssh =~ ^[Yy]$ ]]; then
            print_info "跳过 SSH 配置"
            return
        fi
    fi
    
    # 备份 SSH 配置
    cp /etc/ssh/sshd_config /etc/ssh/sshd_config.bak.$(date +%Y%m%d)
    
    # 优化 SSH 配置
    sed -i 's/#UseDNS yes/UseDNS no/' /etc/ssh/sshd_config
    sed -i 's/GSSAPIAuthentication yes/GSSAPIAuthentication no/' /etc/ssh/sshd_config
    
    if ! grep -q "ClientAliveInterval" /etc/ssh/sshd_config; then
        echo "ClientAliveInterval 60" >> /etc/ssh/sshd_config
        echo "ClientAliveCountMax 3" >> /etc/ssh/sshd_config
    fi
    
    # 重启 SSH 服务
    if command -v systemctl &> /dev/null; then
        systemctl restart sshd 2>/dev/null || systemctl restart ssh 2>/dev/null
    else
        service sshd restart 2>/dev/null || service ssh restart 2>/dev/null
    fi
    
    print_success "SSH 配置优化完成"
}

# 配置时区
configure_timezone() {
    print_info "正在配置时区..."
    
    if [[ -z "${NON_INTERACTIVE:-}" ]]; then
        read -p "是否设置为中国时区 (Asia/Shanghai)? [Y/n]: " set_tz
        set_tz=${set_tz:-Y}
    else
        set_tz="Y"
    fi
    
    if [[ $set_tz =~ ^[Yy]$ ]]; then
        if command -v timedatectl &> /dev/null; then
            timedatectl set-timezone Asia/Shanghai
        else
            ln -sf /usr/share/zoneinfo/Asia/Shanghai /etc/localtime
            echo "Asia/Shanghai" > /etc/timezone 2>/dev/null || true
        fi
        print_success "时区已设置为 Asia/Shanghai"
    fi
}

# 系统优化配置
optimize_system() {
    print_info "正在应用系统优化配置..."
    
    if [[ -z "${NON_INTERACTIVE:-}" ]]; then
        read -p "是否应用系统优化配置? [Y/n]: " do_optimize
        do_optimize=${do_optimize:-Y}
        
        if [[ ! $do_optimize =~ ^[Yy]$ ]]; then
            print_info "跳过系统优化"
            return
        fi
    fi
    
    # 优化文件描述符限制
    if ! grep -q "^* soft nofile 65535" /etc/security/limits.conf; then
        echo "* soft nofile 65535" >> /etc/security/limits.conf
        echo "* hard nofile 65535" >> /etc/security/limits.conf
    fi
    
    # 优化内核参数
    cat > /etc/sysctl.d/99-custom.conf <<EOF
# 网络优化
net.core.rmem_max = 16777216
net.core.wmem_max = 16777216
net.ipv4.tcp_rmem = 4096 87380 16777216
net.ipv4.tcp_wmem = 4096 65536 16777216

# 连接优化
net.ipv4.tcp_max_syn_backlog = 8192
net.core.somaxconn = 8192
net.ipv4.tcp_slow_start_after_idle = 0

# TIME_WAIT 优化
net.ipv4.tcp_fin_timeout = 30
net.ipv4.tcp_tw_reuse = 1

# 文件系统优化
fs.file-max = 1000000
fs.inotify.max_user_watches = 524288
EOF
    
    sysctl -p /etc/sysctl.d/99-custom.conf > /dev/null 2>&1 || true
    
    print_success "系统优化配置完成"
}

# 清理系统
clean_system() {
    print_info "正在清理系统..."
    
    if [[ -z "${NON_INTERACTIVE:-}" ]]; then
        read -p "是否清理系统缓存和无用包? [Y/n]: " do_clean
        do_clean=${do_clean:-Y}
        
        if [[ ! $do_clean =~ ^[Yy]$ ]]; then
            print_info "跳过系统清理"
            return
        fi
    fi
    
    case $OS in
        ubuntu|debian)
            apt-get autoremove -y
            apt-get autoclean -y
            apt-get clean
            ;;
            
        centos|rhel|rocky|almalinux)
            yum autoremove -y
            yum clean all
            ;;
    esac
    
    print_success "系统清理完成"
}

# 显示系统信息
show_system_info() {
    echo ""
    print_success "==================== 系统信息 ===================="
    echo ""
    
    echo "操作系统: $(grep PRETTY_NAME /etc/os-release | cut -d'"' -f2 2>/dev/null || echo "Unknown")"
    echo "内核版本: $(uname -r)"
    echo "系统架构: $(uname -m)"
    
    if command -v nproc &> /dev/null; then
        echo "CPU 核心: $(nproc)"
    fi
    
    if command -v free &> /dev/null; then
        echo "内存大小: $(free -h | awk '/^Mem:/ {print $2}')"
    fi
    
    if command -v df &> /dev/null; then
        echo "根目录空间: $(df -h / | awk 'NR==2 {print $2 " (可用: " $4 ")"}')"
    fi
    
    echo "当前用户: $(whoami)"
    echo "主机名称: $(hostname)"
    
    if [[ -f /etc/timezone ]]; then
        echo "时区设置: $(cat /etc/timezone)"
    elif command -v timedatectl &> /dev/null; then
        echo "时区设置: $(timedatectl | grep "Time zone" | awk '{print $3}')"
    fi
    
    echo ""
    print_success "=================================================="
}

# 显示主菜单
show_menu() {
    echo ""
    echo "请选择要执行的操作:"
    echo ""
    echo "1) 完整初始化 (推荐)"
    echo "   - 更换镜像源"
    echo "   - 安装基础工具"
    echo "   - 系统优化"
    echo "   - SSH 配置"
    echo "   - 时区设置"
    echo "   - 系统清理"
    echo ""
    echo "2) 仅更换镜像源"
    echo "3) 仅安装基础工具"
    echo "4) 仅系统优化"
    echo "5) 仅配置 SSH"
    echo "6) 仅设置时区"
    echo "7) 仅系统清理"
    echo "8) 显示系统信息"
    echo "9) 退出脚本"
    echo ""
}

# 完整初始化流程
full_initialization() {
    print_info "开始完整系统初始化..."
    
    # 1. 备份原始配置
    backup_sources
    
    # 2. 配置镜像源
    case $OS in
        ubuntu|debian)
            configure_ubuntu_mirrors || {
                print_warning "镜像源配置失败，继续其他步骤..."
            }
            ;;
        centos|rhel|rocky|almalinux)
            configure_centos_mirrors || {
                print_warning "镜像源配置失败，继续其他步骤..."
            }
            ;;
        *)
            print_warning "暂不支持该系统的镜像源配置"
            ;;
    esac
    
    # 3. 安装基础工具
    install_basic_tools
    
    # 4. 配置时区
    configure_timezone
    
    # 5. 配置 SSH
    configure_ssh
    
    # 6. 系统优化
    optimize_system
    
    # 7. 系统清理
    clean_system
    
    print_success "系统初始化完成！"
    
    # 显示系统信息
    show_system_info
    
    echo ""
    print_info "建议执行以下操作:"
    echo "1. 重启系统使所有配置生效: reboot"
    echo "2. 检查系统更新: apt-get upgrade 或 yum update"
    echo ""
}

# 显示帮助信息
show_help() {
    echo ""
    echo "Linux 系统初始化脚本 v1.3"
    echo ""
    echo "用法: $0 [选项]"
    echo ""
    echo "选项:"
    echo "  --auto                  自动模式 (非交互式，使用默认配置)"
    echo "  --mirror=SOURCE         指定镜像源 (aliyun, tsinghua, ustc)"
    echo "  --non-interactive       非交互模式 (跳过所有确认)"
    echo "  --help                  显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  sudo $0"
    echo "  sudo $0 --auto"
    echo "  sudo $0 --mirror=tsinghua --non-interactive"
    echo ""
}

# 解析命令行参数
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --auto)
                AUTO_MODE=1
                NON_INTERACTIVE=1
                MIRROR_ARG="aliyun"
                shift
                ;;
            --mirror=*)
                MIRROR_ARG="${1#*=}"
                shift
                ;;
            --non-interactive)
                NON_INTERACTIVE=1
                shift
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                print_error "未知参数: $1"
                show_help
                exit 1
                ;;
        esac
    done
}

# 主函数
main() {
    # 解析命令行参数
    parse_arguments "$@"
    
    # 显示横幅
    print_banner
    
    # 检查权限
    check_permissions
    
    # 检测操作系统
    detect_os
    
    # 自动模式直接执行完整初始化
    if [[ -n "${AUTO_MODE:-}" ]]; then
        full_initialization
        exit 0
    fi
    
    # 交互式菜单
    while true; do
        show_menu
        read -p "请选择操作 [1-9]: " choice
        
        case $choice in
            1)
                full_initialization
                break
                ;;
            2)
                backup_sources
                case $OS in
                    ubuntu|debian)
                        configure_ubuntu_mirrors
                        ;;
                    centos|rhel|rocky|almalinux)
                        configure_centos_mirrors
                        ;;
                esac
                ;;
            3)
                install_basic_tools
                ;;
            4)
                optimize_system
                ;;
            5)
                configure_ssh
                ;;
            6)
                configure_timezone
                ;;
            7)
                clean_system
                ;;
            8)
                show_system_info
                ;;
            9)
                print_info "退出脚本"
                exit 0
                ;;
            *)
                print_warning "无效选择，请重新输入"
                ;;
        esac
        
        # 询问是否继续
        if [[ $choice -ne 9 ]]; then
            echo ""
            read -p "是否继续其他操作? [Y/n]: " continue_choice
            continue_choice=${continue_choice:-Y}
            if [[ ! $continue_choice =~ ^[Yy]$ ]]; then
                print_info "退出脚本"
                exit 0
            fi
        fi
    done
}

# 脚本入口点
main "$@"