#!/usr/bin/env bash
set -euo pipefail

# =============================================
# 环境变量和配置
# =============================================
export GOPATH="${GOPATH:-$HOME/.go/gopath}"
export GOROOT="${GOROOT:-$HOME/.go/goroot}"
export GOROOT_BOOTSTRAP="${GOROOT_BOOTSTRAP:-$HOME/.go/gobootstrap}"
export PATH="$GOPATH/bin:$GOROOT/bin:$PATH"

# =============================================
# 工具函数
# =============================================
# 错误处理
handle_error() {
    local exit_code=$?
    local line=$1
    local command=$2
    echo "Error on line $line: Command '$command' failed with exit code $exit_code"
    exit $exit_code
}

# 命令检查
check_command() {
    if ! command -v "$1" >/dev/null 2>&1; then
        echo "Error: $1 is required but not installed"
        exit 1
    fi
}

# 环境检查
check_env() {
    if [[ -z "$GOROOT" ]]; then
        echo "Error: GOROOT environment variable is not set"
        exit 1
    fi

    if [[ -z "$GOROOT_BOOTSTRAP" ]]; then
        echo "Error: GOROOT_BOOTSTRAP environment variable is not set"
        exit 1
    fi
}

# =============================================
# 版本管理
# =============================================
# 获取可用的 Go 版本
list_versions() {
    check_command jq
    
    curl -s https://go.dev/dl/\?mode\=json\&include\=all | \
        jq -r '.[].version' | \
        grep -v -e 'rc' -e 'beta' | \
        sort -V
}

# 打印版本信息
print_versions() {
    echo "Available Go versions:"
    echo "---------------------"
    
    local versions
    versions=$(list_versions)
    
    local current_version
    if command -v go >/dev/null 2>&1; then
        current_version=$(go version | awk '{print $3}')
    fi
    
    local bootstrap_version
    if [ -x "$GOROOT_BOOTSTRAP/bin/go" ]; then
        bootstrap_version=$("$GOROOT_BOOTSTRAP/bin/go" version | awk '{print $3}')
    fi
    
    while IFS= read -r version; do
        echo "  $version"
    done <<< "$versions"
    
    if [ -n "$current_version" ]; then
        echo "* $current_version (current version)"
    fi
    if [ -n "$bootstrap_version" ]; then
        echo "+ $bootstrap_version (bootstrap version)"
    fi
}

# =============================================
# 安装和升级
# =============================================
# 安装指定版本的 bootstrap
install_bootstrap_version() {
    local version=$1
    check_command curl
    check_command wget
    check_command tar

    local os
    os=$(uname -s | tr '[:upper:]' '[:lower:]')
    local arch
    arch=$(uname -m)
    case "$arch" in
        x86_64) arch="amd64" ;;
        aarch64) arch="arm64" ;;
        *) echo "Unsupported architecture: $arch"; exit 1 ;;
    esac

    local filename="${version}.${os}-${arch}.tar.gz"
    local url="https://go.dev/dl/$filename"
    echo "Downloading Go bootstrap version $version"
    echo "URL: $url"

    mkdir -p "$HOME/.go"
    cd "$HOME/.go" || exit 1
    
    if ! wget -4 "$url"; then
        echo "Failed to download Go bootstrap"
        exit 1
    fi
    
    if ! tar -zxvf "$filename"; then
        echo "Failed to extract Go bootstrap"
        exit 1
    fi
    
    rm -rf "$GOROOT_BOOTSTRAP"
    mv "go" "$version"
    ln -sf "$HOME/.go/$version" "$GOROOT_BOOTSTRAP"
    rm "$filename"
}

# 安装 bootstrap
gobootstrap() {
    local version=$1
    if [ -n "$version" ]; then
        install_bootstrap_version "$version"
        return
    fi

    version=$(curl -s https://go.dev/dl/\?mode\=json | grep version | head -n 1 | sed 's/.*"version": "\(.*\)".*/\1/')

    if [ -d "$GOROOT_BOOTSTRAP" ] && [ -x "$GOROOT_BOOTSTRAP/bin/go" ]; then
        local installed_version
        installed_version=$("$GOROOT_BOOTSTRAP/bin/go" version | awk '{print $3}' | sed 's/go//')
        if [ "$installed_version" = "$version" ]; then
            echo "Go bootstrap is already installed. Skipping download."
            return
        else
            echo "Go bootstrap version mismatch. Removing existing installation."
            rm -rf "$GOROOT_BOOTSTRAP"
        fi
    fi

    install_bootstrap_version "$version"
}

# 升级 Go
upgrade_go() {
    local version=$1
    check_env
    echo "Upgrading Go to version $version"

    cd "$GOROOT" || exit 1
    git fetch --all || exit 1
    
    if [ -n "$version" ]; then
        if ! git tag --list | grep -q "^$version$"; then
            echo "Error: Version $version does not exist"
            exit 1
        fi
        
        local current_version
        if command -v go >/dev/null 2>&1; then
            current_version=$(go version | awk '{print $3}')
        fi
        
        if [ "$version" = "$current_version" ]; then
            echo "$version is already installed"
            exit 0
        fi
    else
        version=$(git tag --list | grep -E '^go[0-9]+\.[0-9]+(\.[0-9]+)?$' | sort -V | tail -n 1)
    fi
    
    version=${version%.*}
    git fetch --all
    git checkout "release-branch.$version" || exit 1
    cd "$GOROOT/src" || exit 1
    ./make.bash || exit 1
}

# 设置 Go root
goroot() {
    local version=$1
    if [ -x "$GOROOT/bin/go" ]; then
        echo "Go binary already exists at $GOROOT/bin/go"
        echo "Current version: $($GOROOT/bin/go version)"
        echo "If you want to reinstall, please remove the $HOME/.go directory first"
        exit 0
    fi

    if [ -d "$HOME/.go/goroot" ] && [ -d "$HOME/.go/goroot/.git" ]; then
        echo "Updating Go root"
        cd "$HOME/.go/goroot"
        git fetch --all
    else
        echo "Setting up Go root"
        if [ -d "$HOME/.go/goroot" ]; then
            echo "Removing existing Go root"
            rm -rf "$HOME/.go/goroot"
        fi

        git clone https://github.com/golang/go.git "$HOME/.go/goroot"
        cd "$HOME/.go/goroot"
    fi

    upgrade_go "$version"
}

# =============================================
# 环境更新
# =============================================
update_env() {
    echo "Updating environment variables"
    echo "
export GOPATH=\"$HOME/.go/gopath\"
export GOROOT_BOOTSTRAP=\"$HOME/.go/gobootstrap\"
export GOROOT=\"$HOME/.go/goroot\"
export PATH=\"$GOPATH/bin:$GOROOT/bin:\$PATH\"
"
}

# =============================================
# 主函数
# =============================================
install_go() {
    local version=$1
    if command -v go >/dev/null 2>&1; then
        echo "Go is already installed at $(which go)"
        echo "Current version: $(go version)"
        echo "If you want to upgrade, please use: $0 upgrade [version]"
        exit 0
    fi

    if [ -d "$GOROOT" ] && [ -d "$GOROOT_BOOTSTRAP" ] && [ -x "$GOROOT_BOOTSTRAP/bin/go" ]; then
        echo "Go environment seems to be already set up at $GOROOT"
        echo "If you want to reinstall, please remove the $HOME/.go directory first"
        exit 0
    fi

    gobootstrap "$version"
    goroot "$version"
    update_env
}

# 安装特定版本的 bootstrap
install_specific_bootstrap() {
    local version=$1
    if [ -z "$version" ]; then
        echo "Error: Please specify a version"
        print_help
        exit 1
    fi

    # 检查版本是否存在
    if ! list_versions | grep -q "^$version$"; then
        echo "Error: Version $version does not exist"
        exit 1
    fi

    install_bootstrap_version "$version"
}

print_help() {
    echo "Usage: $0 [command] [options]"
    echo "Commands:"
    echo "  install [version]     Install Go environment with optional version"
    echo "  upgrade [version]     Upgrade Go to the latest or specified version"
    echo "  list                 List available Go versions"
    echo "  bootstrap [version]  Install specific version of bootstrap"
    echo "  -h, --help          Show help information"
}

# =============================================
# 主程序
# =============================================
trap 'handle_error ${LINENO} "$BASH_COMMAND"' ERR

# 检查必要命令
check_command git
check_command wget
check_command curl
check_command tar

# 处理命令行参数
if [ $# -eq 0 ] || [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
    print_help
    exit 0
fi

case "$1" in
    install)
        install_go "$2"
        ;;
    upgrade)
        upgrade_go "$2"
        ;;
    list)
        print_versions
        ;;
    bootstrap)
        install_specific_bootstrap "$2"
        ;;
    *)
        echo "Unknown command: $1"
        print_help
        exit 1
        ;;
esac