#!/usr/bin/env bash

#
# Setup.
#
export PYTHON_MIRROR=https://mirrors.huaweicloud.com/python/
VERSION="1.0.52"
UP=$'\033[A'
DOWN=$'\033[B'
P_PREFIX=${P_PREFIX-$HOME/.local}
BASE_VERSIONS_DIR=$P_PREFIX/p/versions
#configure flags eg: --enable-optimizations
PY_INSTALL_FLAG=${PY_INSTALL_FLAG-}

insbin(){
  local owner_repo='indygreg/python-build-standalone'; latest_version_url="$(curl -s https://raw.githubusercontent.com/$owner_repo/latest-release/latest-release.json| grep "asset_url_prefix" | cut -d : -f 2,3 | tr -d \")";
  local latest=$(basename $latest_version_url)
  case $(uname | tr '[:upper:]' '[:lower:]') in
    darwin*)    osname="apple-darwin"    ;;
    msys*)    osname=""    ;;
    *)    osname="unknown-linux-gnu" ;;
  esac
  local ver=${1#v}
  cpython=cpython-${ver}+${latest}-x86_64-${osname}-pgo+lto-full.tar.zst
  latest_version_url=$latest_version_url/$cpython
  log fetch "$latest_version_url"
  binhome=${2-$HOME/.local/p/versions/python/$ver}
  test -d $binhome || mkdir -p $binhome
  is_ok "$latest_version_url" || abort "not found version $ver"
  curl -fL $latest_version_url | zstd -d | tar -C $binhome --strip-components=2 -xf - python/install
  export PYTHONHOME=$binhome
}

shell() {
  export PATH=$BASE_VERSIONS_DIR/python:$HOME/.local/bin:$PATH
  get_current_version
  export PATH=$BASE_VERSIONS_DIR/python/$current/bin:$PATH
  bash --rcfile <(echo "export PS1='(py$current)[\w]> '")
}
#
# Add to PATH.
#
ensurepath() {
  BASE_DIR=$BASE_VERSIONS_DIR/python
  if echo $BASE_VERSIONS_DIR | egrep -q ^$HOME; then 
    BASE_DIR=${BASE_DIR/$HOME/"\$HOME"}
  fi
  if test -f $HOME/.bashrc && ! grep -q "export PATH=\"$BASE_DIR:" $HOME/.bashrc; then
      echo -e "\nexport PATH=\"$BASE_DIR:\$HOME/.local/bin:\$PATH\"" >> $HOME/.bashrc
  fi

  if test -f $HOME/.zshrc && ! grep -q "export PATH=\"$BASE_DIR:" $HOME/.zshrc; then
      echo -e "\nexport PATH=\"$BASE_DIR:\$HOME/.local/bin:\$PATH\"" >> $HOME/.zshrc
  fi
  exec $SHELL -l
}

# autoremove unused pkgs
unpkg() {
  test -z "$2" && abort "pip package name required"
  export PATH=$BASE_VERSIONS_DIR/python:$PATH
  python -m pip freeze |grep pip-autoremove || python -m pip install pip-autoremove
  get_current_version
  export PATH=$BASE_VERSIONS_DIR/python/$current/bin:$PATH
  pip-autoremove $2
}

# 
# migrate older package 
#
migration_failed() {
  { echo
    echo "MIGRATION FAILED"
    echo "Inspect or clean up the package list at ${PIP_REQUIREMENTS}"
  } 1>&2
  exit 1
}
migrate() {
  src="$2"
  dst="$3"
  (test -z "$2" || test -z "$3") && abort "version(s) required"

  test -d "$BASE_VERSIONS_DIR"/python"/$src"  || {
    echo "p: not an installed version: $src" 1>&2
    exit 1
  }

  test -d "$BASE_VERSIONS_DIR"/python"/$dst"  || {
    echo "p: not an installed version: $dst" 1>&2
    exit 1
  }
  SEED="$(date "+%Y%m%d%H%M%S").$$"
  PIP_REQUIREMENTS="/tmp/requirements.${SEED}.txt"

  trap migration_failed ERR
  "$BASE_VERSIONS_DIR/python/${src}/bin/python3" -m   pip freeze >> "${PIP_REQUIREMENTS}"
  "$BASE_VERSIONS_DIR/python/${dst}/bin/python3" -m   pip install --requirement "${PIP_REQUIREMENTS}"
  rm -f "${PIP_REQUIREMENTS}"
  trap - ERR
}

#
# Python source.
#

MIRROR=(${PYTHON_MIRROR-https://www.python.org/ftp/python/})

test -d "$BASE_VERSIONS_DIR"/python || mkdir -p "$BASE_VERSIONS_DIR"/python

#
# Log <type> <msg>
#

log() {
    printf "  \033[36m%10s\033[0m : \033[90m%s\033[0m\n" "$1" "$2"
}

#
# Exit with the given <msg ...>
#

abort() {
    printf "\n  \033[31mError: %s\033[0m\n\n" "$@" && exit 1
}

#
# Print a (differentiable from log) success message.
#

success() {
    printf "\n  \033[32mSuccess: %s\033[0m\n\n" "$@"
}

#
# Ensure we have curl or wget support.
#

GET=
command -v curl > /dev/null && GET="curl -# -L"
test -z "$GET" && abort "curl required"

# wget support (Added --no-check-certificate for Github downloads)
#command -v wget > /dev/null && GET="wget --no-check-certificate -q -O-"

#
# Functions used when showing versions installed
#

enter_fullscreen() {
    tput smcup
    stty -echo
}

leave_fullscreen() {
    tput rmcup
    stty echo
}

handle_sigint() {
    leave_fullscreen
    exit $?
}

handle_sigtstp() {
    leave_fullscreen
    kill -s SIGSTOP $$
}

#
# Output usage information.
#

display_help() {
  cat <<-EOF

  Usage: p [COMMAND] [args]

  Commands:

    p                              Output versions installed
    p status                       Output current status
    p <version>                    Activate to Python <version>
      p latest                     Activate to the latest Python release
      p stable                     Activate to the latest stable Python release
    p use <version> [args ...]     Execute Python <version> with [args ...]
    p bin <version>                Output bin path for <version>
    p rm <version ...>             Remove the given version(s)
    p prev                         Revert to the previously activated version
    p default                      Revert to the system default python version
    p ensurepath                   Add python path to your search path
    p unpkg <pkg ...>              Remove the given name(s) pip packages and unused dependencies
    p shell                        Add python bin path to a new bash shell  
    p p                            Wrap pip command
    p px                           Wrap pipx command
    p py                           Wrap poetry command
    p up                           Update pipx all libs,Wrap pipdate command,update all pip packages
    p migrate <ver1> <ver2>        Migrate pip package from a Python version to another
    p ls                           Output the versions of Python available
      p ls latest                  Output the latest Python version available
      p ls stable                  Output the latest stable Python version available
    p i <version>                  Install Python <version> from prebuilt binaries

  Options:

    -V, --version   Output current version of p
    -h, --help      Display help information

  Environment variables:

    P_PREFIX        Python versions installation dir ,default: ~/.local
    PY_INSTALL_FLAG Python installation configure flags, eg: --enable-optimizations
EOF
    exit 0
}

#
# Output status.
#

display_status() {
    get_current_version

    log version "$current"
    log bin "$(display_bin_path_for_version)"
    log previous "$(get_previous_version)"
    log latest "$(is_latest_version)"
    log stable "$(is_latest_stable_version)"
}

#
# Hide cursor.
#

hide_cursor() {
    printf "\e[?25l"
}

#
# Show cursor.
#

show_cursor() {
    printf "\e[?25h"
}

#
# Output version after selected.
#

next_version_installed() {
    list_versions_installed | grep "$selected" -A 1 | tail -n 1
}

#
# Output version before selected.
#

prev_version_installed() {
    list_versions_installed | grep "$selected" -B 1 | head -n 1
}

#
# Output previous version.
#
get_previous_version() {
    if [ ! -f "$BASE_VERSIONS_DIR"/.prev ]; then
        echo "none"
    else
        cat "$BASE_VERSIONS_DIR"/.prev
    fi
}

#
# Output p version.
#

display_p_version() {
    echo $VERSION && exit 0
}

#
# Gets current human-readable Python version.
#
get_current_version() {
    # current=$(python -c 'import sys; print(".".join(map(str, sys.version_info[:3])))')
    local version
    python="$BASE_VERSIONS_DIR"/python/python
    if [[ ! -e "$python" ]] ; then
      python=python
    fi
    version=$($python -V 2>&1)
    current=${version#*Python }
}

#
# Check for installed version, and populate $active
#

check_current_version() {
    #command -v python &> /dev/null
    if test $? -eq 0; then
        get_current_version
        python="$BASE_VERSIONS_DIR"/python/python
        if [[ ! -e "$python" ]] ; then
          python=python
        fi
        if diff &> /dev/null \
            "$BASE_VERSIONS_DIR"/python/"$current"/bin/python3 \
            "$(which $python)" ; then
            active="python/$current"
        fi
    fi
}

#
# Display sorted versions directories paths.
#

versions_paths() {
    find "$BASE_VERSIONS_DIR" -maxdepth 2 -type d \
        | sed 's|'"$BASE_VERSIONS_DIR"'/||g' \
        | egrep "[0-9]+\.[0-9]+\.[0-9]+([a|b]?)([0-9]?)+" \
        | sort -k 1,1n -k 2,2n -k 3,3n -t . -k 4,4n -d -k 5,5n -r
}

#
# Display installed versions with <selected>
#

display_versions_with_selected() {
    selected=$1
    echo
    for version in $(versions_paths); do
        version_minus_python=${version#*python\/}
        if test "$version" = "$selected"; then
            printf "  \033[36mο\033[0m %s\033[0m\n" "$version_minus_python"
        else
            printf "    \033[90m%s\033[0m\n" "$version_minus_python"
        fi
    done
    echo
}

#
# List installed versions.
#

list_versions_installed() {
    for version in $(versions_paths); do
        echo "${version}"
    done
}

#
# Display current python --version and others installed.
#

display_versions() {
    enter_fullscreen
    check_current_version
    display_versions_with_selected "$active"

    trap handle_sigint INT
    trap handle_sigtstp SIGTSTP

    while true; do
        read -r -n 3 c
        case "$c" in
            $UP)
            clear
            display_versions_with_selected "$(prev_version_installed)"
            ;;
            $DOWN)
            clear
            display_versions_with_selected "$(next_version_installed)"
            ;;
            *)
            log activating $selected 
            activate "$selected" leave
            leave_fullscreen

            if [[ "$active" = "" ]]; then
              exec $SHELL -l
            fi
            exit
            ;;
        esac
    done
}

#
# Move up a line and erase.
#

erase_line() {
    printf "\033[1A\033[2K"
}

#
# Check if the HEAD response of <url> is 200.
#

is_ok() {
    curl -Is $1 | head -n 1 | grep -e 200 -e 302 > /dev/null
}

#
# Determine tarball url for <version>
#

tarball_url() {
    version_directory="${version%a*}"
    version_directory="${version_directory%b*}"
    echo "${MIRROR[@]}${version_directory%r*}/Python-${version}.tgz"
}

#
# Activate <version>
#

activate() {
    local version=$1
    check_current_version
    if test "$version" != "$active"; then
        
        test -z "$active" || echo "$active" > "$BASE_VERSIONS_DIR"/.prev

        ln -sf "$BASE_VERSIONS_DIR"/"$version"/bin/python3 "$BASE_VERSIONS_DIR"/python/python
        ln -sf "$BASE_VERSIONS_DIR"/"$version"/bin/python3 "$BASE_VERSIONS_DIR"/python/python3

        if test -d "$HOME/.local/pipx/shared/bin"; then
          log reinstall pipx-libs
          ln -sf "$BASE_VERSIONS_DIR"/"$version"/bin/pip3 $HOME/.local/pipx/shared/bin/pip
          "$BASE_VERSIONS_DIR"/python/python -m pipx reinstall-all &> /dev/null

          if [[ "$active" = "" ]] && [[ "$2" != "leave" ]]; then
            exec $SHELL -l
          fi
        fi
    fi
}

#
# Activate previous Python.
#

activate_previous() {
    test -f "$BASE_VERSIONS_DIR"/.prev || abort "no previous versions activated"
    local prev
    prev=$(cat "$BASE_VERSIONS_DIR"/.prev)
    test -d "$BASE_VERSIONS_DIR"/"$prev" || abort "previous version $prev not installed"
    test ! -z "$prev" || abort "previous version $prev not installed"
    activate "$prev"
    echo
    get_current_version
    log activate "$current"
    echo

    success "Now using Python $current!"
}

#
# Activate default (prior to p) Python.
#
activate_default() {
    log activate default

    rm -f "$BASE_VERSIONS_DIR"/python/python
    rm -f "$BASE_VERSIONS_DIR"/python/python3
    exec $SHELL -l

    get_current_version
    success "Now using system default Python $current!\n  Use \`p <version>\` to activate another version."
}

#
# Install <version>
#
OPENSSL=$P_PREFIX/ssl

installdeps() {
  if [ ! -d "$OPENSSL" ]; then
    command -v sudo > /dev/null && sudo=sudo
    $sudo apt-get install -y build-essential zlib1g-dev libffi-dev zstd &>/dev/null
    $sudo yum install -y which diffutils perl gcc gcc-c++ zlib-devel libffi-devel zstd &>/dev/null

    log installing openssl
    curl -L -o /tmp/openssl.tar.gz https://www.openssl.org/source/openssl-1.1.1t.tar.gz
    cd /tmp && tar xf openssl*.gz && cd openssl-*
    ./config --prefix=$OPENSSL no-zlib &> /dev/null 
    make -j 8 &> /dev/null
    make install &> /dev/null
  fi
}

install() {
    local version=${1#v}

    local dir=$BASE_VERSIONS_DIR/python/tmp/$version
    local pydir=$BASE_VERSIONS_DIR/python/$version
    local url
    url=$(tarball_url "$version")

    if test -d "$pydir"; then
        if [[ ! -e "$dir"/p.lock ]] ; then
            log activate "$version"

            activate python/"$version"

            get_current_version
            
            success "Now using Python $current!"

            exit
        fi
    fi

    echo

    log installing Python-"$version"

    is_ok "$url" || abort "invalid version $version"

    log create "$dir"
    mkdir -p "$dir"
    if [ $? -ne 0 ] ; then
        abort "sudo required"
    else
        touch "$dir"/p.lock
    fi

    cd "$dir" || exit

    if [[ ! "$2" == "bin" ]]; then
      log installing system-dependencies
      installdeps

      log fetch "$url"
      curl -#L "$url" | tar -zx --strip 1
      erase_line

      log configure "$version"
      ./configure  "$PY_INSTALL_FLAG" --prefix=$pydir  --with-openssl=$OPENSSL  --with-ensurepip=install  &> /dev/null

      log compile "$version"
      make -j 8 &> /dev/null
      make install &> /dev/null

      if [[ ! -f "python" ]] && [[ ! -f "python.exe" ]]; then
        abort "Unable to compile Python $version!"
      fi
    else
      insbin $version $pydir
    fi

    activate python/"$version" leave
    log activate "$version"

    source ~/.bashrc
    test -f ~/.zshrc && source ~/.zshrc
    export PATH=$BASE_VERSIONS_DIR/python:$PATH
    log refresh $PATH
    log  install pipx 
    python -m pip install pipx pipdate
    log  install poetry
    python -m pipx install --force poetry
    mv ~/.local/pipx/shared/pyvenv.cfg /tmp 2>/dev/null

    get_current_version
    rm -rf "$dir"
    success "Now using Python $current!"
    success "Now you can exec \`p shell\` for the temporary use or exec \`p ensurepath\` to add python bin PATH to .bashrc"
}

#
# Remove <version ...>
#

remove_versions() {
    test -z "$1" && abort "version(s) required"

    for version in "$@"; do
        if [[ "$version" = "latest" ]]; then
            version=$(display_latest_version)
        elif [[ "$version" = "stable" ]]; then
            version=$(display_latest_stable_version)
        fi

        get_current_version
        rm -rf "$BASE_VERSIONS_DIR"/python/"${version#v}"
        log remove "$version"

        if [[ "$current" = "${version#v}" ]]; then
          rm -f "$BASE_VERSIONS_DIR"/python/python
          rm -f "$BASE_VERSIONS_DIR"/python/python3
          exec $SHELL -l
        fi
    done

    versions=("$@")

    success "Removed Python ${versions// /, }!"
}

#
# Output bin path for <version>
#

display_bin_path_for_version() {
    get_current_version

    if [ ! -z "$1" ]; then
        local version=${1#v}
    else
        if [ ! -d "$BASE_VERSIONS_DIR"/python/"$current" ]; then
            abort "Version required!"
        else
            local version=$current;
        fi
    fi

    local bin=$BASE_VERSIONS_DIR/python/$version/bin/python3
    if test -f "$bin"; then
        printf "%s \n" "$bin"
    else
        abort "Python $version is not installed"
    fi
}

#
# Execute the given <version> of Python with [args ...]
#

execute_with_version() {
    test -z "$1" && abort "version required"
    local version=${1#v}

    if [[ "$version" = "latest" ]]; then
        version=$(display_latest_version)
    elif [[ "$version" = "stable" ]]; then
        version=$(display_latest_stable_version)
    fi

    local bin=$BASE_VERSIONS_DIR/python/$version/bin/python3

    shift # remove version

    if test -f "$bin"; then
        $bin "$@"
    else
        abort "Python $version is not installed"
    fi
}

#
# Display the latest release version.
#

display_latest_version() {
    latest_directory=$($GET 2> /dev/null "${MIRROR[@]}" \
        | egrep -o '[0-9]+\.[0-9]+\.[0-9]+' \
        | sort -u -k 1,1n -k 2,2n -k 3,3n -t . \
        | tail -n1)

    $GET 2> /dev/null "${MIRROR[@]}$latest_directory/" \
        | egrep -o '[0-9]+\.[0-9]+\.[0-9]+[a|b|r]c?[0-9]+' \
        | sort -k 1,1n -k 2,2n -k 3,3n -t . -k 4,4n -d -k 5,5n -d \
        | tail -n1
}

#
# Determine if current version is the latest version.
#

is_latest_version() {
    get_current_version
    if [[ $current == $(display_latest_version) ]]; then
        echo "yes"
    else
        echo "no"
    fi
}

#
# Display the latest stable release version.
#

display_latest_stable_version() {
    $GET 2> /dev/null "${MIRROR[@]}" \
        | egrep -o '[0-9]+\.[0-9]*[02468]\.[0-9]+' \
        | sort -u -k 1,1n -k 2,2n -k 3,3n -t . \
        | tail -n1
}

#
# Determine if current version is the latest stable version.
#

is_latest_stable_version() {
    get_current_version
    if [[ $current == $(display_latest_stable_version) ]]; then
        echo "yes"
    else
        echo "no"
    fi
}

#
# Display the versions available.
#

display_remote_versions() {
    check_current_version
    local versions=""
    versions=$($GET 2> /dev/null "${MIRROR[@]}" \
    | egrep -o '[0-9]+\.[0-9]+\.[0-9]+' \
    | egrep -v '^0\.[0-7]\.' \
    | egrep -v '^0\.8\.[0-5]$' \
    | sort -u -k 1,1n -k 2,2n -k 3,3n -t . \
    | awk '{ print "  " $1 }')

    echo

    for v in $versions; do
        if test "$active" = "python/$v"; then
            printf "  \033[36mο\033[0m %s \033[0m\n" "$v"
        else
            if test -d "$BASE_VERSIONS_DIR"/python/"$v"; then
                printf "    %s \033[0m\n" "$v"
            else
                printf "    \033[90m%s\033[0m\n" "$v"
            fi
        fi
    done
    echo
}

#
# Handle arguments.
#
python="$BASE_VERSIONS_DIR"/python/python
checks="test ! -e $python && abort \"no installed version\""

if test $# -eq 0; then
    test -z "$(versions_paths)" && abort "no installed version"
    display_versions
else
    while test $# -ne 0; do
        case $1 in
            -V|--version) display_p_version ;;
            -h|--help|help) display_help ;;
            status) display_status ;;
            bin|which)
            case $2 in
                latest) display_bin_path_for_version "$($0 ls latest)"; exit ;;
                stable) display_bin_path_for_version "$($0 ls stable)"; exit ;;
                *) display_bin_path_for_version "$2"; exit ;;
            esac
            exit ;;
            as|use) shift; execute_with_version "$@"; exit ;;
            rm|-) shift; remove_versions "$@"; exit ;;
            ls|list)
            case $2 in
                latest) display_latest_version; exit ;;
                stable) display_latest_stable_version; exit ;;
                *) display_remote_versions; exit ;;
            esac
            exit ;;
            prev|previous) activate_previous; exit ;;
            ensurepath) ensurepath; exit ;;
            unpkg) unpkg "$@"; exit ;;
            migrate) migrate "$@"; exit ;;
            shell) shell; exit ;;
            p) eval $checks; $python -m pip "${@:2}"; exit ;;
            px) eval $checks; $python -m pipx "${@:2}"; exit ;;
            py) eval $checks; ~/.local/bin/poetry "${@:2}"; exit ;;
            up) eval $checks; $python -m pipx upgrade-all ;
                $python -m pipdate "${@:2}"; exit ;;
            default) activate_default; exit ;;
            latest) install "$($0 ls latest)"; exit ;;
            stable) install "$($0 ls stable)"; exit ;;
            i) install "$2" bin; exit ;;
            *) install "$1"; exit ;;
        esac
        shift
    done
fi
