#!/bin/bash

# =================================================================
# 脚本名称: ubuntu_init.sh
# 描述: Ubuntu 系统初始化工具 (系统信息、Root SSH、Docker 运维、系统更新、环境优化)
# 制作人: kikock
# =================================================================

# 检查是否以 root 权限运行
if [[ $EUID -ne 0 ]]; then
   echo "错误: 请使用 sudo 或 root 用户运行此脚本。"
   exit 1
fi

# 定义颜色常量
GREEN='\033[0;32m'
BLUE='\033[0;34m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m' # 重置颜色

# 1. 查看系统信息函数
show_sys_info() {
    clear
    echo -e "${BLUE}================ 系统详细信息 ================${NC}"

    OS_NAME=$(lsb_release -ds 2>/dev/null || cat /etc/os-release | grep "PRETTY_NAME" | cut -d= -f2 | tr -d '"')
    KERNEL=$(uname -r)
    ARCH=$(uname -m)
	CPU_MODEL=$(lscpu | grep "Model name" | cut -d: -f2 | xargs)

   # 增强版内存抓取：兼容不同版本的 free 输出格式
    # 移除冒号匹配以兼容中文环境的全角冒号 [cite: 5.1]
    MEM_LINE=$(free -h | grep -E "Mem|内存")
    # 如果关键字匹配失败，尝试直接读取第二行作为兜底
    if [ -z "$MEM_LINE" ]; then
        MEM_LINE=$(free -h | sed -n '2p')
    fi

    MEM_TOTAL=$(echo $MEM_LINE | awk '{print $2}')
    MEM_USED=$(echo $MEM_LINE | awk '{print $3}')
    MEM_FREE=$(echo $MEM_LINE | awk '{print $4}')

    SWAP_LINE=$(free -h | grep -E "Swap|交换")
    # 如果关键字匹配失败，尝试直接读取第三行作为兜底
    if [ -z "$SWAP_LINE" ]; then
        SWAP_LINE=$(free -h | sed -n '3p')
    fi

    SWAP_TOTAL=$(echo $SWAP_LINE | awk '{print $2}')
    SWAP_USED=$(echo $SWAP_LINE | awk '{print $3}')


    DISK_USAGE=$(df -h / | awk '/\// {print $5}')
    IP_ADDR=$(hostname -I | cut -d' ' -f1)
    UPTIME=$(uptime -p)

    # 获取 NPU 实时频率 (针对 Kernel 6.1 路径)
    if [ -f "/sys/class/devfreq/fdab0000.npu/cur_freq" ]; then
        NPU_FREQ=$(cat /sys/class/devfreq/fdab0000.npu/cur_freq | awk '{printf "%.0f MHz\n", $1/1000000}')
    else
        NPU_FREQ="未激活/挂起"
    fi

    # 获取 NPU 驱动版本
    if [ -f "/sys/kernel/debug/rknpu/version" ]; then
        NPU_VER=$(cat /sys/kernel/debug/rknpu/version | grep -oE "v[0-9]+\.[0-9]+\.[0-9]+")
    elif lsmod | grep -q "rknpu"; then
        NPU_VER=$(modinfo rknpu 2>/dev/null | grep "^version:" | awk '{print $2}')
        [ -z "$NPU_VER" ] && NPU_VER="已加载(版本未知)"
    else
        NPU_VER="未检测到驱动"
    fi

    # 获取系统温度
    if [ -f "/sys/class/thermal/thermal_zone0/temp" ]; then
        TEMP=$(cat /sys/class/thermal/thermal_zone0/temp | awk '{print $1/1000 "°C"}')
    else
        TEMP="未知"
    fi


    echo -e "${YELLOW}操作系统:${NC}  $OS_NAME"
    echo -e "${YELLOW}内核版本:${NC}  $KERNEL"
    echo -e "${YELLOW}系统架构:${NC}  $ARCH"
    echo -e "${YELLOW}处理器:${NC}    ${CPU_MODEL:-未知}"
    echo -e "${YELLOW}根分区占用:${NC} $DISK_USAGE"
	echo -e "${YELLOW}物理内存:${NC}  $MEM_TOTAL (已用: $MEM_USED / 空闲: $MEM_FREE)"
    echo -e "${YELLOW}虚拟内存:${NC}  $SWAP_TOTAL (已用: $SWAP_USED)"
    echo -e "${YELLOW}NPU 频率:${NC}  $NPU_FREQ"
    echo -e "${YELLOW}NPU 驱动:${NC}  $NPU_VER"
    echo -e "${YELLOW}核心温度:${NC}  $TEMP"
    echo -e "${YELLOW}内网 IP:${NC}   $IP_ADDR"
    echo -e "${YELLOW}运行时间:${NC}  $UPTIME"
    echo -e "${BLUE}==============================================${NC}"
    echo ""
    read -p "按回车键返回主菜单..."
}

# 2. 修改 Root 密码并开启 SSH 登录
setup_root_ssh() {
    echo -e "\n${BLUE}--- 开始配置 Root SSH 登录 ---${NC}"
    echo -e "${YELLOW}请设置新的 Root 密码:${NC}"
    passwd root

    if [ $? -eq 0 ]; then
        echo -e "${BLUE}正在修改 SSH 配置文件...${NC}"
        cp /etc/ssh/sshd_config /etc/ssh/sshd_config.bak.$(date +%F_%T)
        sed -i 's/^#\?PermitRootLogin.*/PermitRootLogin yes/' /etc/ssh/sshd_config
        sed -i 's/^#\?PasswordAuthentication.*/PasswordAuthentication yes/' /etc/ssh/sshd_config
        systemctl restart ssh

        if [ $? -eq 0 ]; then
            echo -e "${GREEN}成功: Root 密码已设置，且已允许 SSH 远程登录。${NC}"
        else
            echo -e "${RED}错误: SSH 服务重启失败。${NC}"
        fi
    else
        echo -e "${RED}取消: 密码修改失败或用户取消操作。${NC}"
    fi
    echo ""
    read -p "按回车键返回主菜单..."
}

# 3. 离线安装 Docker 逻辑 (已优化解压路径)
install_docker_offline() {
    echo -e "\n${BLUE}--- 开始离线安装 Docker ---${NC}"

    if command -v docker &> /dev/null; then
        echo -e "${YELLOW}检测到系统已安装 Docker，脚本将进行覆盖升级/重置操作。${NC}"
    fi

    read -p "请输入 Docker 离线包名称 (如 docker-26.1.3.tgz): " TAR_NAME
    TAR_FILE="./docker/$TAR_NAME"
    TEMP_DIR="./dockerApp"

    if [ ! -f "$TAR_FILE" ]; then
        echo -e "${RED}错误: 在 ./docker/ 目录下找不到文件 $TAR_NAME${NC}"
        return
    fi

    # 环境准备
    echo -e "${YELLOW}正在处理 SELinux ...${NC}"
    if [ -f /etc/selinux/config ]; then
        sed -i 's/SELINUX=enforcing/SELINUX=disabled/' /etc/selinux/config
        setenforce 0 2>/dev/null
    fi

    # 创建隔离的解压目录
    [ -d "$TEMP_DIR" ] && rm -rf "$TEMP_DIR"
    mkdir -p "$TEMP_DIR"

    echo -e "${YELLOW}正在解压 $TAR_NAME 到 $TEMP_DIR ...${NC}"
    tar -xzvf "$TAR_FILE" -C "$TEMP_DIR"
    if [ $? -ne 0 ]; then
        echo -e "${RED}错误: 解压失败。${NC}"
        rm -rf "$TEMP_DIR"
        return
    fi

    # 配置目录与文件分发 (从临时目录拷贝)
    echo -e "${YELLOW}分发二进制文件与配置...${NC}"
    [ ! -d "/etc/docker" ] && mkdir -p "/etc/docker"

    # 拷贝解压后的 etc 目录内容
    if [ -d "$TEMP_DIR/etc" ]; then
        cp -r "$TEMP_DIR/etc/"* /etc/
    fi

    # 拷贝解压后的 docker 二进制文件
    if [ -d "$TEMP_DIR/docker" ]; then
        cp "$TEMP_DIR/docker/"* /usr/bin/
    fi

    # 核心组件来源: ./docker/ (用户指定的下级目录)
    if [ -f "./docker/docker.service" ]; then
        cp ./docker/docker.service /usr/lib/systemd/system/
        chmod +x /usr/lib/systemd/system/docker.service
    else
        echo -e "${RED}警告: 在 ./docker/ 目录下未发现 docker.service 文件。${NC}"
    fi

    echo -e "${YELLOW}正在加载配置并启动/重启 Docker 服务...${NC}"
    systemctl daemon-reload
    # 使用 restart 确保如果是覆盖安装能立即生效
    systemctl restart docker
    systemctl enable docker.service

    # 安装 docker-compose (来源: ./docker/)
    if [ -f "./docker/docker-compose" ]; then
        echo -e "${YELLOW}配置 docker-compose...${NC}"
        cp ./docker/docker-compose /usr/local/bin/
        chmod +x /usr/local/bin/docker-compose
        ln -sf /usr/local/bin/docker-compose /usr/bin/docker-compose
    fi

    # 清理临时解压目录
    echo -e "${YELLOW}清理临时文件...${NC}"
    rm -rf "$TEMP_DIR"

    echo -e "${GREEN}Docker 离线安装/更新任务执行完毕。${NC}"
    docker -v
}

# 4. 卸载 Docker 逻辑
uninstall_docker() {
    echo -e "\n${RED}--- 开始卸载 Docker ---${NC}"
    read -p "确定要彻底卸载 Docker 及其组件吗? (y/n): " confirm
    [[ "$confirm" != "y" ]] && return

    systemctl stop docker 2>/dev/null
    systemctl disable docker 2>/dev/null
    rm -f /usr/bin/docker* /usr/bin/containerd* /usr/bin/runc /usr/bin/ctr
    rm -f /usr/local/bin/docker-compose /usr/bin/docker-compose
    rm -f /usr/lib/systemd/system/docker.service
    rm -rf /etc/docker
    systemctl daemon-reload
    echo -e "${GREEN}Docker 卸载完成。${NC}"
}

# 5. 常用 Docker 命令说明
show_docker_commands() {
    clear
    echo -e "${BLUE}================ 常用 Docker 命令参考 ================${NC}"
    echo -e "${YELLOW}[ 容器管理 ]${NC}"
    echo "  docker ps          - 查看运行中的容器"
    echo "  docker ps -a       - 查看所有容器 (包括已停止)"
    echo "  docker start [ID]  - 启动容器"
    echo "  docker stop [ID]   - 停止容器"
    echo "  docker restart [ID]- 重启容器"
    echo "  docker logs -f [ID]- 查看容器实时日志"
    echo "  docker exec -it [ID] /bin/bash - 进入容器终端"
    echo ""
    echo -e "${YELLOW}[ 镜像管理 ]${NC}"
    echo "  docker images      - 列出本地所有镜像"
    echo "  docker rmi [ID]    - 删除本地镜像"
    echo "  docker pull [NAME] - 从仓库拉取镜像"
    echo ""
    echo -e "${YELLOW}[ Docker Compose ]${NC}"
    echo "  docker-compose up -d  - 后台启动所有服务"
    echo "  docker-compose down   - 停止并移除所有服务"
    echo "  docker-compose ps     - 查看服务状态"
    echo -e "${BLUE}======================================================${NC}"
    echo ""
    read -p "按回车键返回..."
}

# 6. 安装 Docker Web 管理功能 (docker-fast)
install_docker_fast_web() {
    echo -e "\n${BLUE}--- 开始安装 Docker Web 管理终端 ---${NC}"

    # 检查 Docker 是否安装
    if ! command -v docker &> /dev/null; then
        echo -e "${RED}错误: 未检测到 Docker，请先安装 Docker。${NC}"
        return
    fi

    # 检查 docker-fast.yml 是否存在
    YML_FILE="./docker/docker-fast.yml"
    if [ ! -f "$YML_FILE" ]; then
        echo -e "${RED}错误: 在 ./docker/ 目录下找不到 docker-fast.yml 文件。${NC}"
        return
    fi

    # 检查是否已经运行
    if docker ps --format '{{.Names}}' | grep -q "docker-fast"; then
        echo -e "${YELLOW}检测到 Docker-Fast 已在运行，将进行配置检查并尝试更新。${NC}"
    fi

    echo -e "${YELLOW}正在使用 docker-compose 部署/更新 Web 管理端...${NC}"

    # 优先尝试 docker-compose 命令，其次尝试 docker compose
    if command -v docker-compose &> /dev/null; then
        docker-compose -f "$YML_FILE" up -d
    elif docker compose version &> /dev/null; then
        docker compose -f "$YML_FILE" up -d
    else
        echo -e "${RED}错误: 未检测到 docker-compose，请检查安装情况。${NC}"
        return
    fi

    if [ $? -eq 0 ]; then
        echo -e "${GREEN}成功: Docker Web 管理端已启动/更新。${NC}"
        echo -e "${BLUE}提示: 请检查 $YML_FILE 中的端口映射，并使用浏览器访问服务器 IP。${NC}"
    else
        echo -e "${RED}错误: 部署失败。${NC}"
    fi
}

# 7. Docker 管理二级菜单
docker_menu() {
    while true; do
        clear
        echo -e "${GREEN}==============================================${NC}"
        echo -e "${GREEN}             Docker 管理 (二级菜单)           ${NC}"
        echo -e "${GREEN}==============================================${NC}"
        echo " 1. 离线安装 Docker"
        echo " 2. 彻底卸载 Docker"
        echo " 3. 常用 Docker 命令说明"
        echo " 4. 安装 Docker Web 管理终端 (docker-fast)"
        echo " 0. 返回主菜单"
        echo -e "${GREEN}==============================================${NC}"
        read -p "请选择操作 [0-4]: " sub_choice

        case $sub_choice in
            1) install_docker_offline; read -p "按回车键继续..." ;;
            2) uninstall_docker; read -p "按回车键继续..." ;;
            3) show_docker_commands ;;
            4) install_docker_fast_web; read -p "按回车键继续..." ;;
            0) break ;;
            *) echo -e "${RED}无效输入。${NC}"; sleep 1 ;;
        esac
    done
}

# 8. 系统软件包更新函数
update_system_packages() {
    echo -e "\n${BLUE}--- 开始更新系统软件包 ---${NC}"
    echo -e "${YELLOW}1. 正在更新软件源 (apt update)...${NC}"
    apt update
    if [ $? -eq 0 ]; then
        echo -e "${YELLOW}2. 正在升级软件包 (apt upgrade)...${NC}"
        apt upgrade -y
        if [ $? -eq 0 ]; then
            echo -e "${GREEN}成功: 系统软件包已全部更新完毕。${NC}"
        else
            echo -e "${RED}错误: 软件包升级过程中出现问题。${NC}"
        fi
    else
        echo -e "${RED}错误: 无法更新软件源，请检查网络。${NC}"
    fi
    echo ""
    read -p "按回车键返回主菜单..."
}

system_optimization_menu() {
    # 定义颜色常量 (如果外部未定义)
    local GREEN='\033[0;32m'
    local YELLOW='\033[1;33m'
    local RED='\033[0;31m'
    local BLUE='\033[0;34m'
    local NC='\033[0m'

    # Root 权限检查
    if [[ $EUID -ne 0 ]]; then
       echo -e "${RED}错误: 必须以 root 权限运行此脚本。${NC}"
       return 1
    fi

    while true; do
        clear
        echo -e "${GREEN}==============================================${NC}"
        echo -e "${GREEN}            系统环境优化 (二级菜单)            ${NC}"
        echo -e "${GREEN}==============================================${NC}"
        echo -e " 1. 切换至阿里云镜像源 (Ubuntu 专用)"
        echo -e " 2. 开启内核 BBR 网络加速"
        echo -e " 3. 配置 2GB zRAM 虚拟内存 (高压缩比)"
        echo -e " 4. 清除 zRAM/Swap 所有相关配置"
        echo -e " 5. 设置系统时区为上海 (Asia/Shanghai)"
        echo -e " 0. 返回主菜单"
        echo -e "${GREEN}==============================================${NC}"
        read -p "请选择操作 [0-5]: " opt_choice

        case $opt_choice in
            1)
                echo -e "${YELLOW}正在备份并尝试切换镜像源...${NC}"
                cp /etc/apt/sources.list /etc/apt/sources.list.bak 2>/dev/null
                # 兼容旧版 sources.list 和新版 sources.list.d 结构 (Ubuntu 24.04+)
                sed -i 's|http://.*.ubuntu.com|http://mirrors.aliyun.com|g' /etc/apt/sources.list
                apt update
                echo -e "${GREEN}成功: 镜像源处理完成。${NC}"
                read -p "按回车键继续..."
                ;;
            2)
                echo -e "${YELLOW}正在配置 BBR 加速...${NC}"
                if ! grep -q "net.core.default_qdisc=fq" /etc/sysctl.conf; then
                    echo "net.core.default_qdisc=fq" >> /etc/sysctl.conf
                    echo "net.ipv4.tcp_congestion_control=bbr" >> /etc/sysctl.conf
                    sysctl -p
                    echo -e "${GREEN}成功: BBR 已开启。${NC}"
                else
                    echo -e "${BLUE}提示: BBR 似乎已经开启，无需重复操作。${NC}"
                fi
                read -p "按回车键继续..."
                ;;
            3)
                echo -e "${YELLOW}正在部署 zRAM 2GB 持久化方案...${NC}"
                if swapon --show | grep -q "/dev/zram0"; then
                    echo -e "${BLUE}提示: zRAM 已经在运行中。${NC}"
                else
                    # 创建初始化脚本
                    cat > /usr/local/bin/zram-init.sh <<EOF
#!/bin/bash
/sbin/modprobe zram num_devices=1
echo lzo > /sys/block/zram0/comp_algorithm
echo 2147483648 > /sys/block/zram0/disksize
/sbin/mkswap /dev/zram0
/sbin/swapon /dev/zram0 -p 10
EOF
                    chmod +x /usr/local/bin/zram-init.sh

                    # 创建 Systemd 服务
                    cat > /etc/systemd/system/zram-setup.service <<EOF
[Unit]
Description=Setup zRAM Swap
After=multi-user.target

[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/local/bin/zram-init.sh

[Install]
WantedBy=multi-user.target
EOF
                    systemctl daemon-reload
                    systemctl enable --now zram-setup.service
                    echo -e "${GREEN}成功: 2GB zRAM 已开启并设置为开机自启。${NC}"
                fi
                sleep 2
                ;;
            4)
                echo -e "${YELLOW}正在彻底移除虚拟内存配置...${NC}"
                systemctl stop zram-setup.service 2>/dev/null
                systemctl disable zram-setup.service 2>/dev/null
                rm -f /etc/systemd/system/zram-setup.service
                rm -f /usr/local/bin/zram-init.sh
                swapoff /dev/zram0 2>/dev/null
                modprobe -r zram 2>/dev/null
                systemctl daemon-reload
                echo -e "${GREEN}成功: 配置已清除。${NC}"
                sleep 2
                ;;
            5)
                echo -e "${YELLOW}正在更新系统时区...${NC}"
                timedatectl set-timezone Asia/Shanghai
                echo -e "${GREEN}成功: 当前系统时间: $(date)${NC}"
                read -p "按回车键继续..."
                ;;
            0)
                echo -e "${BLUE}返回中...${NC}"
                break
                ;;
            *)
                echo -e "${RED}输入有误，请输入 0-5 之间的数字。${NC}"
                sleep 1
                ;;
        esac
    done
}

# 主菜单循环
while true; do
    clear
    echo -e "${GREEN}==============================================${NC}"
    echo -e "${GREEN}      Ubuntu 系统初始化工具箱 - By kikock     ${NC}"
    echo -e "${GREEN}==============================================${NC}"
    echo " 1. 查看系统详细信息 (包含架构)"
    echo " 2. 修改 Root 密码并开启 SSH 登录"
    echo " 3. Docker 管理 (安装/卸载/命令/Web)"
    echo " 4. 系统软件包更新 (apt update & upgrade)"
    echo " 5. 系统环境优化 (源/BBR/Swap/时区)"
    echo " 0. 退出脚本"
    echo -e "${GREEN}==============================================${NC}"
    read -p "请输入选项 [0-5]: " choice

    case $choice in
        1) show_sys_info ;;
        2) setup_root_ssh ;;
        3) docker_menu ;;
        4) update_system_packages ;;
        5) system_optimization_menu ;;
        0) echo -e "${BLUE}感谢使用，再见！- kikock${NC}"; exit 0 ;;
        *) echo -e "${RED}输入无效，请重新选择。${NC}" ; sleep 1 ;;
    esac
done