#!/bin/bash

is_pipe() {
    [ -t 0 ] && return 1 || return 0
}

# 设置颜色变量
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
NC='\033[0m'  # No Color

# 确保脚本以root权限运行
check_root() {
  if [ "$(id -u)" != "0" ]; then
    echo -e "${RED}此脚本需要root权限${NC}" 1>&2
    exit 1
  fi
}

# 更新软件源
update_sources() {
  echo -e "${YELLOW}正在更新软件源，请稍候...${NC}"
  apt update -y
  echo -e "${GREEN}软件源更新完成！${NC}"
}

# 升级所有包，保留本地修改的配置文件
upgrade_packages() {
  echo -e "${YELLOW}正在升级所有软件包，请稍候...${NC}"
  apt upgrade -y --option=Dpkg::Options::="--force-confold"
  echo -e "${GREEN}软件包升级完成！${NC}"
}

# 安装常用工具和软件包
install_tools() {
  echo -e "${YELLOW}正在安装常用工具和软件包，请稍候...${NC}"
  apt install wget curl sudo vim git sl lrzsz -y
  mv /usr/games/sl /usr/bin/sl
  echo "alias sl='sl -e'" >> /etc/profile
  echo -e "${GREEN}常用工具和软件包安装完成！${NC}"
}

# 备份并去除python增强
remove_python_enhancements() {
  echo -e "${YELLOW}正在备份并去除 Python 增强功能，请稍候...${NC}"
  PYTHON_LIB_DIR=$(ls /usr/lib/ | grep -E '^python3\.[0-13]+$' | sort -V | tail -n 1)
  cp /usr/lib/$PYTHON_LIB_DIR/EXTERNALLY-MANAGED /usr/lib/$PYTHON_LIB_DIR/EXTERNALLY-MANAGED.bk
  rm /usr/lib/$PYTHON_LIB_DIR/EXTERNALLY-MANAGED
  echo -e "${GREEN}Python 增强功能已移除！${NC}"
}

# 修改vim默认配置
configure_vim() {
  echo -e "${YELLOW}正在修改 Vim 默认配置，请稍候...${NC}"
  find /usr/share/vim -name "defaults.vim" -exec sed -i 's/mouse=a/mouse=""/g' {} \;
  echo -e "${GREEN}Vim 默认配置已修改！${NC}"
}

# 去掉/root/.bashrc中的#号
clean_bashrc() {
  echo -e "${YELLOW}正在清理 /root/.bashrc 文件，请稍候...${NC}"
  sed -i -r '/^#[[:space:]]*(alias|export|umask|eval)/ s/^#//' /root/.bashrc
  echo -e "${GREEN}/root/.bashrc 文件清理完成！${NC}"
}

# 允许root用户远程登录，并允许密码登录
configure_ssh() {
  echo -e "${YELLOW}正在配置 SSH，请稍候...${NC}"
  sed -i 's/#\?PermitRootLogin .*/PermitRootLogin yes/' /etc/ssh/sshd_config
  sed -i 's/#\?PasswordAuthentication .*/PasswordAuthentication yes/' /etc/ssh/sshd_config
  systemctl restart sshd
  echo -e "${GREEN}SSH 配置完成！${NC}"
}

# 检查并设置时区为Asia/Shanghai
check_and_set_timezone() {
  echo -e "${YELLOW}正在检查并设置时区，请稍候...${NC}"
  current_timezone=$(timedatectl show --property=Timezone --value)
  if [ "$current_timezone" != "Asia/Shanghai" ]; then
    if timedatectl set-timezone Asia/Shanghai 2>/dev/null; then
      echo -e "${GREEN}时区已修改为 Asia/Shanghai。${NC}"
    else
      echo -e "${RED}错误：无法修改时区。${NC}"
    fi
  else
    echo -e "${GREEN}时区已经是 Asia/Shanghai。${NC}"
  fi
}

# 检查、安装和配置NTP服务
setup_ntp() {
  echo -e "${YELLOW}正在检查、安装和配置 NTP 服务，请稍候...${NC}"
  if ! systemctl is-active --quiet ntpsec; then
    echo -e "${YELLOW}需要安装 NTP 服务，请稍等${NC}"
    if apt-get update && apt-get install -y --option=Dpkg::Options::="--force-confold" ntpsec ; then
      echo -e "${GREEN}NTPsec 服务已安装。${NC}"
    else
      echo -e "${RED}错误：NTPsec 服务安装失败。${NC}"
      return 1
    fi
  fi

  # 判断服务器位置并配置服务器
  if ping -c 1 google.com &> /dev/null; then
    echo -e "${YELLOW}服务器可能位于国外，配置国际NTP服务器${NC}"
    servers=(
      "server time.cloudflare.com iburst"
      "server time.facebook.com iburst"
      "server time.google.com iburst"
    )
  else
    echo -e "${YELLOW}服务器可能位于中国，配置中国NTP服务器${NC}"
    servers=(
      "server ntp.aliyun.com iburst"
      "server ntp2.tencent.com iburst"
    )
  fi

  # 生成新的 NTP 配置文件
  cat > /etc/ntpsec/ntp.conf << EOF
# /etc/ntpsec/ntp.conf, configuration for ntpd; see ntp.conf(5) for help
driftfile /var/lib/ntpsec/ntp.drift
leapfile /usr/share/zoneinfo/leap-seconds.list
tos maxclock 11
tos minclock 4 minsane 3
$(printf "%s\n" "${servers[@]}")
restrict default kod nomodify nopeer noquery limited
restrict 127.0.0.1
restrict ::1
EOF

  echo -e "${GREEN}NTP 配置文件已更新。${NC}"
}

# 重启NTP服务并检查时间同步状态
restart_and_check_ntp() {
  echo -e "${YELLOW}正在重启 NTP 服务并检查时间同步状态，请稍候...${NC}"
  if systemctl restart ntpsec; then
    echo -e "${GREEN}NTPsec 服务已重启。${NC}"
  else
    echo -e "${RED}错误：无法重启 NTPsec 服务。${NC}"
    return 1
  fi

  # 等待一段时间让NTP服务同步
  sleep 10

  # 使用awk分析ntpq -pn的输出，检查offset字段
  if ntpq -pn | awk '
    BEGIN { sync = 0 }
    {
      if ($1 ~ /^\*/) { sync = 1 }  # 如果有*标记的行，设置sync为1
      if ($9 && $9 < 100 && $9 > -100) { sync = 1 }  # 检查offset值是否在-100到100毫秒之间
    }
    END { exit !sync }'  # 如果sync为1，退出状态为0（成功），否则为1（失败）
  then
    echo -e "${GREEN}NTP 时间同步正常。${NC}"
  else
    echo -e "${RED}NTP 时间同步可能有问题，请手动检查。${NC}"
  fi
}

# 处理 dpkg 被中断的错误
handle_dpkg_interrupted() {
  if [[ -f /var/lib/dpkg/status-old ]]; then
    echo -e "${YELLOW}检测到 dpkg 被中断，正在尝试修复...${NC}"
    sudo dpkg --configure -a || true  # 尝试修复，即使失败也继续执行
    echo -e "${GREEN}dpkg 修复完成！${NC}"
    # 删除标记文件，避免重复触发
    rm -f /var/lib/dpkg/status-old
  fi
}

# 安装 Python 3
install_python3() {
  echo -e "${YELLOW}正在安装 Python 3，请稍候...${NC}"
  apt install -y python3 python3-pip
  echo -e "${GREEN}Python 3 安装完成！${NC}"
}

# 安装 uv
install_uv() {
  echo -e "${YELLOW}正在安装 uv，请稍候...${NC}"
  curl -LsSf https://astral.sh/uv/install.sh | sh
  echo -e "${GREEN}uv 安装完成！${NC}"
}

# 交互式函数
interactive_choices() {
  local disable_ipv6=false
  local install_uv=true  # 默认安装 uv

  if ! is_pipe; then
    read -p "是否禁用 IPv6？(y/n): " choice
    [[ $choice =~ ^[Yy]$ ]] && disable_ipv6=true

    read -p "是否安装 uv (Python 替代品)？默认是 (y/n): " choice
    [[ $choice =~ ^[Nn]$ ]] && install_uv=false
  else
    echo "脚本通过管道执行，使用默认选项：不禁用 IPv6，安装 uv"
  fi

  if $disable_ipv6; then
    echo -e "${YELLOW}正在禁用 IPv6，请稍候...${NC}"
    echo "net.ipv6.conf.all.disable_ipv6 = 1" >> /etc/sysctl.conf && sysctl -p
    echo -e "${GREEN}IPv6 已禁用！${NC}"
  fi

  if $install_uv; then
    install_uv
  else
    install_python3
  fi
}


# 主函数
main() {
  check_root
  handle_dpkg_interrupted
  update_sources
  upgrade_packages
  install_tools
  remove_python_enhancements
  configure_vim
  clean_bashrc
  configure_ssh
  check_and_set_timezone
  setup_ntp
  restart_and_check_ntp
  interactive_choices
  echo -e "${GREEN}初始化完成，请退出终端重新连接${NC}"
}

# 执行主函数
main

