#!/bin/bash

# 检查是否以 root 权限运行 (用于系统级安装)
check_root() {
  if [ "$EUID" -ne 0 ]; then
    echo "注意：选择系统级安装 (/usr/local) 或清理某些旧路径可能需要 root 权限。"
  fi
}

# 提示用户选择安装方式 (最新或指定版本)
select_version() {
  echo "=== Go 语言安装脚本 ==="
  if [ -t 0 ]; then # 如果是交互模式
    echo "请选择安装方式："
    echo "1) 安装最新版本的 Go (推荐)"
    echo "2) 指定安装某一版本的 Go"
    read -p "请输入选项 (1 或 2, 默认: 1): " version_choice
  fi
  version_choice=${version_choice:-1} # 默认选择 1

  local arch="linux-amd64" # 可以根据实际情况检测架构
  GO_ARCH="$arch"

  if [ "$version_choice" == "2" ]; then
    if [ -t 0 ]; then
        read -p "请输入要安装的 Go 版本号 (例如 1.22.2): " version
    else
        echo "非交互模式下，无法指定版本，将安装最新版本。"
        version_choice=1 # 退回到安装最新版本
    fi
    if [ "$version_choice" == "2" ]; then # 再次检查，如果仍然是指定版本
        if [ -z "$version" ]; then
          echo "未指定版本号，退出。"
          exit 1
        fi
        GO_VERSION_NAME="go${version}.${GO_ARCH}"
        DOWNLOAD_URL="https://go.dev/dl/${GO_VERSION_NAME}.tar.gz"
        GO_VERSION_FILENAME="${GO_VERSION_NAME}.tar.gz"
    fi
  fi

  # 如果是选择1或者从2退回
  if [ "$version_choice" == "1" ]; then
    echo "正在获取 Go 语言最新版本下载链接..."
    LATEST_FILENAME=$(curl -s https://go.dev/dl/ | grep -oP 'go[0-9.]+\.'${GO_ARCH}'\.tar\.gz' | head -1)
    if [ -z "$LATEST_FILENAME" ]; then
      echo "错误：未能获取到最新版本 Go 的下载链接，请检查网络或稍后再试。"
      exit 1
    fi
    DOWNLOAD_URL="https://go.dev/dl/$LATEST_FILENAME"
    GO_VERSION_FILENAME="$LATEST_FILENAME"
    echo "检测到最新版本: ${LATEST_FILENAME%.tar.gz}"
  fi

  echo "将下载: $DOWNLOAD_URL"
}

# 提示用户选择安装路径
select_install_path() {
  echo -e "\n--- 选择 Go 语言安装路径 ---"
  if [ ! -t 0 ]; then # 如果 stdin 不是一个终端 (非交互模式)
    echo "检测到非交互模式，将使用默认安装路径: $HOME/go"
    path_choice=2 # 强制选择用户级安装
  else
    echo "1) 系统级安装至 /usr/local (推荐用于多用户环境, 可能需要 root 权限)"
    echo "2) 用户级安装至 $HOME/go (推荐用于个人开发环境)"
    echo "3) 自定义安装路径"
    read -p "请输入选项 (1, 2 或 3, 默认: 2): " path_choice
    path_choice=${path_choice:-2} # 默认选择 2
  fi

  INSTALL_ROOT=""
  GO_INSTALL_DIR=""

  case "$path_choice" in
    1)
      INSTALL_ROOT="/usr/local"
      GO_INSTALL_DIR="/usr/local/go"
      if [ "$EUID" -ne 0 ]; then
          echo "警告：选择系统级安装 (/usr/local) 通常需要 root 权限。如果后续操作失败，请尝试使用 sudo 运行脚本。"
      fi
      ;;
    2)
      INSTALL_ROOT="$HOME"
      GO_INSTALL_DIR="$HOME/go"
      ;;
    3)
      if [ ! -t 0 ]; then
          echo "错误：自定义安装路径在非交互模式下不可用。将使用默认路径 $HOME/go。"
          INSTALL_ROOT="$HOME"
          GO_INSTALL_DIR="$HOME/go"
      else
          read -p "请输入自定义安装路径 (例如 /opt/go): " custom_path
          if [ -z "$custom_path" ]; then
            echo "未指定自定义路径，退出。"
            exit 1
          fi
          # 确保自定义路径是绝对路径
          if [[ "$custom_path" != /* ]]; then
              echo "错误：自定义路径必须是绝对路径 (以 / 开头)。"
              exit 1
          fi
          INSTALL_ROOT=$(dirname "$custom_path")
          GO_INSTALL_DIR="$custom_path"
          echo "将安装至: $GO_INSTALL_DIR"
          # 尝试创建父目录
          mkdir -p "$INSTALL_ROOT"
          if [ ! -d "$INSTALL_ROOT" ] || [ ! -w "$INSTALL_ROOT" ]; then
             echo "错误：自定义路径的父目录 '$INSTALL_ROOT' 无法创建或不可写。"
             exit 1
          fi
      fi
      ;;
    *)
      echo "无效选项，退出。"
      exit 1
      ;;
  esac

  echo "Go 将被安装到: $GO_INSTALL_DIR"
}

# 查找并移除旧的 Go 安装
remove_old_go() {
  echo -e "\n--- 查找并处理旧的 Go 安装 ---"
  local old_go_path_in_env=""
  local old_goroot_inferred=""

  # 1. 尝试通过 PATH 查找现有的 go 命令
  if command -v go &>/dev/null; then
    old_go_path_in_env=$(command -v go)
    echo "检测到当前 PATH 中存在 Go 命令: $old_go_path_in_env"
    old_goroot_inferred=$(dirname "$(dirname "$old_go_path_in_env")")

    if [ -d "$old_goroot_inferred/bin" ] && [ -f "$old_goroot_inferred/VERSION" ]; then
      echo "推断旧的 Go 安装根目录 (GOROOT) 可能为: $old_goroot_inferred"

      if [ "$old_goroot_inferred" == "$GO_INSTALL_DIR" ]; then
        echo "新安装路径与推断的旧 Go 安装路径相同 ($GO_INSTALL_DIR)。将先移除它。"
        if [ -d "$GO_INSTALL_DIR" ]; then
          echo "正在移除: $GO_INSTALL_DIR ..."
          rm -rf "$GO_INSTALL_DIR"
          if [ $? -eq 0 ]; then echo "路径 $GO_INSTALL_DIR 已成功移除。"; else echo "警告：移除 $GO_INSTALL_DIR 失败。"; fi
        fi
        echo "--- 旧 Go 安装处理完毕 (与新路径相同) ---"
        return # 已处理与新路径相同的旧路径，提前返回
      else
        if [ -t 0 ]; then
          read -p "是否要移除这个推断的旧 Go 安装目录 '$old_goroot_inferred'? (y/n, 默认 n): " remove_choice
          remove_choice=${remove_choice:-n}
          if [[ "$remove_choice" =~ ^[Yy]$ ]]; then
            echo "正在移除旧的 Go 安装: $old_goroot_inferred ..."
            rm -rf "$old_goroot_inferred"
            if [ $? -eq 0 ]; then echo "旧的 Go 安装 $old_goroot_inferred 已成功移除。"; else echo "警告：移除 $old_goroot_inferred 失败。"; fi
          else
            echo "跳过移除推断的旧 Go 安装 $old_goroot_inferred。"
          fi
        else
          echo "非交互模式，跳过移除推断的旧 Go 安装 $old_goroot_inferred (因其与新安装路径不同)。"
        fi
      fi
    else
      echo "推断的路径 $old_goroot_inferred 看起来不像一个有效的 Go 安装目录，跳过处理。"
    fi
  else
    echo "在当前 PATH 中未检测到 'go' 命令。"
  fi

  # 2. 检查并移除新安装路径下可能存在的旧文件 (保底逻辑)
  echo "正在检查新安装目标路径 $GO_INSTALL_DIR 是否有残留文件..."
  if [ -d "$GO_INSTALL_DIR" ]; then
    # 只有当上面的逻辑没有移除这个目录时才执行
    # (即 old_goroot_inferred 为空，或者 old_goroot_inferred 不是 GO_INSTALL_DIR)
    if [ -z "$old_goroot_inferred" ] || [ "$old_goroot_inferred" != "$GO_INSTALL_DIR" ]; then
        if [ -t 0 ]; then
            read -p "新安装路径 '$GO_INSTALL_DIR' 已存在。是否要清空它以进行全新安装? (y/n, 默认 y): " clear_target_choice
            clear_target_choice=${clear_target_choice:-y}
        else
            clear_target_choice="y" # 非交互模式默认清空
        fi

        if [[ "$clear_target_choice" =~ ^[Yy]$ ]]; then
            echo "正在清空目标安装路径: $GO_INSTALL_DIR ..."
            rm -rf "$GO_INSTALL_DIR"
            if [ $? -eq 0 ]; then echo "目标安装路径 $GO_INSTALL_DIR 已清空。"; else echo "警告：清空目标安装路径 $GO_INSTALL_DIR 失败。"; fi
        else
            echo "跳过清空目标安装路径 $GO_INSTALL_DIR。注意：这可能导致文件冲突。"
        fi
    fi
  else
    echo "新安装目标路径 $GO_INSTALL_DIR 不存在，无需清理。"
  fi
  echo "--- 旧 Go 安装处理完毕 ---"
}


# 下载并解压 Go
download_and_extract() {
  echo -e "\n--- 下载并解压 Go ---"
  echo "正在下载 $DOWNLOAD_URL ..."
  local temp_go_tar="/tmp/${GO_VERSION_FILENAME}"
  wget --tries=3 -q "$DOWNLOAD_URL" -O "$temp_go_tar"
  if [ ! -f "$temp_go_tar" ] || [ ! -s "$temp_go_tar" ]; then # 检查文件存在且大小不为0
    echo "错误：Go 安装包下载失败或文件为空。"
    exit 1
  fi
  echo "下载成功，正在解压至 $INSTALL_ROOT ..."
  # 确保 INSTALL_ROOT 存在 (特别是自定义路径时，父目录可能需要被创建)
  mkdir -p "$INSTALL_ROOT"
  tar -C "$INSTALL_ROOT" -xzf "$temp_go_tar"
  if [ $? -ne 0 ]; then
     echo "错误：Go 安装包解压失败。"
     rm -f "$temp_go_tar"
     exit 1
  fi
  echo "解压完成。"
  rm -f "$temp_go_tar"
}

# 配置环境变量
configure_environment() {
  echo -e "\n--- 配置环境变量 ---"
  local go_bin_path="$GO_INSTALL_DIR/bin"
  # 用户的 GOPATH 通常是 $HOME/go，确保它存在
  local user_gopath="$HOME/go"
  mkdir -p "$user_gopath/bin" "$user_gopath/pkg" "$user_gopath/src"


  local profile_files_to_check=("$HOME/.bashrc" "$HOME/.profile" "$HOME/.zshrc")
  local user_profile_configured=false

  for profile in "${profile_files_to_check[@]}"; do
    if [ -f "$profile" ]; then
      local needs_goroot=true
      local needs_gopath=true
      local needs_path=true

      # 检查 GOROOT
      if grep -q "export GOROOT=[\"']$GO_INSTALL_DIR[\"']" "$profile"; then
        needs_goroot=false
      fi
      # 检查 GOPATH
      if grep -q "export GOPATH=[\"']$user_gopath[\"']" "$profile"; then
        needs_gopath=false
      fi
      # 检查 PATH
      if grep -q "export PATH=.*[\"']$go_bin_path[\"']" "$profile"; then
        needs_path=false
      fi

      if [ "$needs_goroot" = true ] || [ "$needs_gopath" = true ] || [ "$needs_path" = true ]; then
        echo "正在为 $profile 添加/更新 Go 环境变量..."
        # 如果文件末尾没有空行，先加一个，避免直接追加到最后一行
        [[ $(tail -c1 "$profile" | wc -l) -eq 0 ]] && echo >> "$profile"
        echo -e "\n# Go language environment variables (added by install script)" >> "$profile"

        if [ "$needs_goroot" = true ]; then
          echo "export GOROOT=\"$GO_INSTALL_DIR\"" >> "$profile"
        fi
        if [ "$needs_gopath" = true ]; then
          echo "export GOPATH=\"$user_gopath\"" >> "$profile"
        fi
        if [ "$needs_path" = true ]; then
          # 更安全的 PATH 添加方式，避免重复
          echo "export PATH=\$PATH:\"$go_bin_path\"" >> "$profile"
          # 如果 GOPATH/bin 也想加入 PATH (常见做法)
          # echo "export PATH=\$PATH:\"$go_bin_path\":\"$user_gopath/bin\"" >> "$profile"
        fi
        echo "Go 环境变量已写入 $profile"
        user_profile_configured=true
      else
        echo "Go 环境变量似乎已在 $profile 中正确配置。"
      fi
    fi
  done

  if [ "$user_profile_configured" = true ]; then
    echo "用户级环境变量已配置。请执行 'source ~/.bashrc' 或 'source ~/.profile' (根据您使用的 shell) 或重启终端使其生效。"
  else
     echo "未找到常见的用户配置文件或配置已存在。"
     echo "如果 Go 命令不可用，请确认以下内容已在您的 shell 配置文件中 (例如 ~/.bashrc):"
     echo "export GOROOT=\"$GO_INSTALL_DIR\""
     echo "export GOPATH=\"$user_gopath\""
     echo "export PATH=\$PATH:\"$go_bin_path\""
  fi

  # 系统级配置选项 (仅在 root 且交互模式下询问)
  if [ "$EUID" -eq 0 ] && [ -t 0 ]; then
    echo -e "\n是否同时配置系统级环境变量 (/etc/profile)? 这会影响所有用户。"
    read -p "请输入 (y/n, 默认: n): " system_profile_choice
    system_profile_choice=${system_profile_choice:-n}

    if [[ "$system_profile_choice" =~ ^[Yy]$ ]]; then
      if ! grep -q "export GOROOT=[\"']$GO_INSTALL_DIR[\"']" /etc/profile; then
        echo "正在为 /etc/profile 添加 Go 环境变量..."
        # 如果文件末尾没有空行
        [[ $(tail -c1 /etc/profile | wc -l) -eq 0 ]] && echo >> /etc/profile
        echo -e "\n# Go language system-wide environment variables (added by install script)" >> /etc/profile
        echo "export GOROOT=\"$GO_INSTALL_DIR\"" >> /etc/profile
        # 系统级的 GOPATH 通常不指向特定用户的 HOME，这里可以不设置或设置一个公共路径
        # echo "export GOPATH=\"/opt/go_projects\"" >> /etc/profile
        echo "export PATH=\$PATH:\"$go_bin_path\"" >> /etc/profile
        echo "Go 环境变量已写入 /etc/profile。"
        echo "请注意：/etc/profile 通常在登录时加载，可能需要重启系统或登录会话才能完全生效。"
      else
        echo "Go GOROOT 似乎已在 /etc/profile 中配置，跳过。"
      fi
    fi
  fi
  echo "--- 环境变量配置完毕 ---"
}

# 验证安装
verify_installation() {
  echo -e "\n--- 验证 Go 安装 ---"
  local go_executable="$GO_INSTALL_DIR/bin/go"

  if [ ! -x "$go_executable" ]; then
      echo "错误: 无法找到或执行 Go 命令在 $go_executable。"
      echo "请检查安装路径和环境变量配置。"
      return 1 # 表示验证失败
  fi

  echo "使用 $go_executable version 进行验证:"
  # 设置临时的 GOROOT 和 PATH 来执行验证，避免依赖外部环境
  local temp_goroot="$GO_INSTALL_DIR"
  local temp_path_prefix="$GO_INSTALL_DIR/bin"
  local current_path="$PATH" # 保存当前PATH

  # 尝试运行 go version
  GOROOT="$temp_goroot" PATH="$temp_path_prefix:$current_path" "$go_executable" version
  local exit_code=$?
  local version_string=$("$go_executable" version 2>/dev/null) # 获取版本字符串，忽略可能的 GOROOT/GOPATH 警告

  if [ $exit_code -eq 0 ]; then
    echo "Go 安装成功，版本为: $version_string"
    echo "请打开一个新的终端窗口，或运行 'source ~/.bashrc' (或您的shell对应的配置文件) 来使用 Go 命令。"
  else
    echo "Go 安装可能失败或验证时出现问题。 '$go_executable version' 返回错误码 $exit_code。"
    echo "请检查 $GO_INSTALL_DIR 目录是否完整，以及环境变量是否正确配置并生效。"
  fi
  return $exit_code
}

# --- 脚本主流程 ---
main() {
  check_root
  select_version
  select_install_path
  remove_old_go
  download_and_extract
  configure_environment
  verify_installation

  if [ $? -eq 0 ]; then
    echo -e "\n=== Go 安装成功完成 ==="
  else
    echo -e "\n=== Go 安装过程中出现问题，请检查上面的日志 ==="
  fi
}

# 执行主函数
main

