#!/bin/sh

# This is a little script that can be downloaded from the internet to install the Atlas CLI.
# It attempts to determine architecture and distribution to download and install the correct binary.
# It runs on most Unix shells like {a,ba,da,k,z}sh.

set -u

C_VERSION="0.3.0"
C_FLAVOR_COMMUNITY="community"
C_DEFAULT_ATLAS_INSTALL_PATH="/usr/local/bin/atlas"

ATLAS_NO_UPDATE_NOTIFIER=true
ATLAS_UPDATE_SRV="${ATLAS_UPDATE_SRV:-https://atlasbinaries.com/atlas}"
ATLAS_UPDATE_SRV_FALLBACK="https://release.ariga.io/atlas"
ATLAS_VERSION="${ATLAS_VERSION:-latest}"
ATLAS_FLAVOR="${ATLAS_FLAVOR:-}"
ATLAS_DEBUG="${ATLAS_DEBUG:-false}"
CI="${CI:-false}"

# Alias "beta" to "extended"
if [ "$ATLAS_VERSION" = "beta" ]; then
    ATLAS_VERSION="extended-latest"
fi

# If ATLAS_FLAVOR is set to a beta-only driver, use extended version
if [ -n "$ATLAS_FLAVOR" ]; then
    case "$ATLAS_FLAVOR" in
        snowflake|spanner|databricks|oracle)
            if [ "$ATLAS_VERSION" = "latest" ]; then
                ATLAS_VERSION="extended-latest"
            else
                ATLAS_VERSION="extended-$ATLAS_VERSION"
            fi
            ;;
        *)
            err "Unknown ATLAS_FLAVOR: $ATLAS_FLAVOR. Supported: oracle, snowflake, spanner and databricks"
            ;;
    esac
fi

main() {
    # Prerequisites
    need_cmd uname
    need_cmd mktemp
    need_cmd chmod
    need_cmd chown
    need_cmd mkdir

    local _skip_prompt=false
    local _install=true
    local _path=$C_DEFAULT_ATLAS_INSTALL_PATH
    local _path_set=false
    local _community=false
    local _user_platform _ua_extend
    # Consume flags
    while [ $# -gt 0 ]; do
        case "$1" in
        "-y" | "--yes")
            _skip_prompt=true
            ;;
        "-o" | "--output")
            shift
            _path=$1
            _path_set=true
            ;;
        "--community")
            _community=true
            ;;
        "--platform")
            shift
            _user_platform=$1
            ;;
        "--no-install")
            _install=false
            ;;
        "--user-agent")
            shift
            _ua_extend=$1
            ;;
        "--help")
            printHelp
            ;;
        esac
        shift
    done

    # If the binary should only be downloaded but not installed, path must be set.
    if [ "$_install" = false ] && [ "$_path_set" = false ]; then
        err "Flag '--no-install' requires flag '-o | --output' to be set as well."
    fi

    # Under some circumstances we skip prompting,
    # e.g. in CI systems or if there is no tty to read input from.
    if [ "$CI" = "true" ] || ! sh -c ': >/dev/tty' >/dev/null 2>&1; then
        _skip_prompt=true
    fi

    if [ "$_path_set" = false ] && check_cmd "atlas"; then
        local _out="$(atlas version 2>/dev/null)"
        # Ensure this binary is ours before attempting to override.
        # Run 'atlas version' and ensure it contains the string 'ariga'.
        # If so, proceed. If not, ask where to save (if possible).
        if [ "${_out#*/ariga/atlas/}" != "$_out" ]; then
            _path="$(which atlas)"
        elif [ "$_skip_prompt" = false ]; then
            echo "Found '$_path', but it is not Atlas CLI."
            echo "To override '$_path' type 'yes'."
            echo "To install to a different location, specify the path."
            local _prompt="Type 'yes' or path: "
            if [ -t 0 ]; then
                read -p "$_prompt" _path
            else
                read -p "$_prompt" _path </dev/tty
            fi
            if [ "$_path" = "yes" ]; then
                _path=$C_DEFAULT_ATLAS_INSTALL_PATH
            fi
            if [ -z "$_path" ]; then
                err "Invalid path: '$_path'"
            fi
        else
            # Either there is no tty or prompting was disabled by flag.
            err "Found '$_path', but it is not Atlas CLI.\nSpecify installation location using the '-o PATH' flag!"
        fi
    fi

    # Determine architecture and cpu type.
    get_architecture || return 1
    local _platform="${_user_platform:-$PLATFORM}"
    local _os="$OS"
    local _chipset="$CHIPSET"

    # Create a temporary download folder.
    local _dir
    _dir="$(ensure mktemp -d)"
    ensure mkdir -p "$_dir"

    # Build filename and download path.
    local _file="atlas"
    if [ "$_community" = true ]; then
        _file="$_file-$C_FLAVOR_COMMUNITY"
    fi
    _file="$_file-$_platform-$ATLAS_VERSION"
    case "$_platform" in *windows*)
        _file="$_file.exe"
    esac
    # Build list of servers to try (primary + fallback if set)
    local _servers="$ATLAS_UPDATE_SRV"
    if [ -n "$ATLAS_UPDATE_SRV_FALLBACK" ]; then
        _servers="$_servers $ATLAS_UPDATE_SRV_FALLBACK"
    fi

    # Try each server until one succeeds
    local _url _success=false _head_response_status
    for _server in $_servers; do
        _url="$_server/$_file"
        if [ "$ATLAS_DEBUG" = "true" ]; then
            _url="$_url?test=1"
        fi

        # Check if the binary exists on this server
        _head_response_status=$(curl --head --silent --write-out '%{http_code}' --output /dev/null "$_url" 2>/dev/null || echo "000")
        if [ "$_head_response_status" -eq 404 ] && [ "$ATLAS_VERSION" != "latest" ]; then
            err "The specified Atlas version does not exist or is no longer supported.\\nSee: https://atlasgo.io/cli-reference#distributed-binaries"
        fi
        if [ "$_head_response_status" -eq 200 ] || [ "$_head_response_status" -eq 302 ] || [ "$_head_response_status" -eq 301 ]; then
            _success=true
            break
        fi
    done

    if [ "$_success" = false ]; then
        err "Failed to find binary on any configured server. Tried: $_servers"
    fi

    # Since this script is most likely piped into sh, there is no stdin to prompt the user on.
    # In that case explicitly connect to /dev/tty to read user input.
    if [ "$_skip_prompt" = false ]; then
        # We want to prompt the user. If there is a stdin, use it. If not, use /dev/tty.
        local _yn _prompt
        if [ "$_install" = true ]; then
            _prompt="Install"
        else
            _prompt="Download"
        fi
        local _prompt="$_prompt '$_file' to '$_path'? [y/N] "
        if [ -t 0 ]; then
            read -p "$_prompt" _yn
        else
            read -p "$_prompt" _yn </dev/tty
        fi
        case "$_yn" in
        "y" | "yes" | "Y" | "YES") ;;
        *) exit ;;
        esac
    fi

    # Try downloading from each server until one succeeds
    _success=false
    local _curlVersion="$(curl --version | head -n 1 | awk '{ print $2 }')"
    for _server in $_servers; do
        _url="$_server/$_file"
        if [ "$ATLAS_DEBUG" = "true" ]; then
            _url="$_url?test=1"
        fi

        echo "Downloading $_url"
        local _ua="Atlas Installer/$C_VERSION ($_os; $_chipset) cURL/$_curlVersion"
        if [ -n "${_ua_extend-}" ]; then
            _ua="$_ua $_ua_extend"
        fi
        cd "$_dir"
        if curl -L -o "$_file" -A "$_ua" --fail "-#" "$_url"; then
            _success=true
            break
        else
            echo "Failed to download from $_server, trying next server..." >&2
        fi
    done

    if [ "$_success" = false ]; then
        err "Failed to download binary from any configured server. Tried: $_servers"
    fi
    back

    # If the user wanted us to install the binary for him do so.
    if [ "$_install" = true ]; then
        case "$_platform" in
        *linux*)
            # Install the binary in path
            local _install="install -o root -g root -m 0755 $_dir/$_file $_path"
            if ! [ "$(id -u)" = 0 ]; then
                _install="sudo $_install"
            fi
            ensure_silent eval "$_install"
            ;;
        *darwin*)
            ensure chmod +x "$_dir/$_file"
            # On Mac sometimes the default path does not exist.
            local _mkdir="mkdir -p ${_path%/*}"
            local _mv="mv $_dir/$_file $_path"
            if ! [ "$(id -u)" = 0 ]; then
                _mkdir="sudo $_mkdir"
                _mv="sudo $_mv"
            fi
            ensure_silent eval "$_mkdir"
            ensure_silent eval "$_mv"
            if ! [ "$(id -u)" = 0 ]; then
                ensure_silent sudo chown root: "$_path"
            fi
            ;;
        esac
        # Run once to ensure atlas is installed correctly.
        ensure "$_path" version
        echo "Installation successful!"
    else
        mkdir -p ${_path%/*}
        mv "$_dir/$_file" "$_path"
        echo "Download successful!"
    fi
}

get_architecture() {
    local _ostype _cputype _os
    _ostype="$(uname -s)"
    _cputype="$(uname -m)"

    if [ "$_ostype" = Darwin ] && [ "$_cputype" = i386 ]; then
        # Darwin `uname -m` lies
        if sysctl hw.optional.x86_64 | grep -q ': 1'; then
            _cputype=x86_64
        fi
    fi

    CHIPSET="$_cputype"

    case "$_cputype" in
    xscale | arm | armv6l | armv7l | armv8l | aarch64 | arm64)
        _cputype=arm64
        ;;
    x86_64 | x86-64 | x64 | amd64)
        _cputype=amd64
        ;;
    *)
        err "unknown CPU type: $_cputype"
        ;;
    esac

    case "$_ostype" in
    Linux | FreeBSD | NetBSD | DragonFly)
        _ostype=linux
        _os=Linux
        ;;
    Darwin)
        _ostype=darwin
        _os=MacOS
        ;;
    *)
        err "unrecognized OS type: $_ostype"
        ;;
    esac

    PLATFORM="$_ostype-$_cputype"
    OS="$_os"
}

need_cmd() {
    if ! check_cmd "$1"; then
        err "need '$1' (command not found)"
    fi
}

check_cmd() {
    command -v "$1" >/dev/null 2>&1
}

ensure() {
    if ! "$@"; then err "command failed: $*"; fi
}

ensure_silent() {
    if ! "$@"; then exit 1; fi
}

err() {
    echo "$1" >&2
    exit 1
}

version() {
    echo "$@" | awk -F. '{ printf("%d%03d%03d%03d\n", $1,$2,$3,$4); }'
}

back() {
    >/dev/null cd -
}

printHelp() {
    echo "Atlas CLI Shell Installer.

Usage:
    curl -sSf https://atlasgo.sh | sh
    curl -sSf https://atlasgo.sh | sh -s -- [options]

Options:
    -y, --yes     Skip interactive prompting defaulting to 'yes'
    -o, --output  Path to save the binary to, attempting to create it, if it does not exist
    --community   Install the community version
    --no-install  Download only without adding to PATH
    --platform    Specify platform to assume, e.g. 'darwin-arm64', 'linux-amd64', ...
    --user-agent  Add comments to User Agent information sent by cURL request
    --help        Print this help message

Environment:
    ATLAS_VERSION  Set the version to install (default: 'latest')
    ATLAS_FLAVOR   Set to 'oracle', 'snowflake', 'spanner', or 'databricks' to install the extended version with that driver
    CI             Disables prompting for CI systems
"
    exit 0
}

main "$@" || exit 1