#!/usr/bin/env zsh
# Simple script which automatically defines certain aliases for python,
# which will automatically use certain python version 
# Versions are automatically obtained from $PYENV_ROOT/versions directories
# NOTE: This assumes that all folders in this directory are valid python versions
#
# Assume we have these installed pyenv python versions 3.6.5, 3.6.12 and 3.10.1:
#   - Set full-version aliases: py3.6.5, py3.6.12 and py3.10.1
#   - Set py3 to 3.10.1 (latest with major version 3)
#   - Set py3.6 to 3.6.12 (latest with major version 3 and minor version 6)
#   - Set py3.10 to 3.10.1 (latest, and only python with major version 3 and minor version 10)


# Define all wanted aliases for a given python version
# $1 - full valid pyenv python version (for example '3.6.12', `3.11-dev`, or `pypy3.6-7.2.0-src`)
# $2 - version used in the alias (for example '3.6', '3', or even '', but also `pypy3.7`, ...)
define_aliases() {
    version="$1"
    alias_version="$2"
    cmd_prefix="PYENV_VERSION=$version"

    alias "py$alias_version=$cmd_prefix python"
    alias "ipy$alias_version=$cmd_prefix ipython"
    alias "bpy$alias_version=$cmd_prefix bpython"
    alias "pydoc$alias_version=$cmd_prefix pydoc"
    alias "pytest$alias_version=$cmd_prefix pytest"
    alias -g "PY$alias_version=$cmd_prefix"
}

# Handle splitting full version into prefix, version number and suffix
# Because of the huge variaty of python implemenations and their different namings,
# this function will only be able to handle the default CPython version names,
# which follow the regex pattern of: '\d+\.\d+\.\d+', the rest will print 'full_version;;'
# In the future, this may also include support for some other naming schemes.
# $1 - full valid pyenv python version (for example '3.6.12', `3.11-dev`, or `pypy3.6-7.2.0-src`)
parse_python_version() {
    full_version="$1"
    if echo "$full_version" | grep -E "[0-9]+\.[0-9]+\.[0-9]+" >/dev/null; then
        echo ";$full_version;"
    else
        echo ';;'
    fi
}

# Prints version number extracted from alias for given version
# $1 - version used in the alias (for example '3.6', '3', or even '', but also 'pypy3.6', ...)
get_alias_version() {
    alias_version="$1"
    definition="$(alias "py$alias_version")"
    full_version="$(echo "$definition" | cut -d= -f3 | cut -d' '  -f1)"

    version_info="$(parse_python_version "$full_version")"
    version="$(echo "$version_info" | cut -d';' -f2)"
    echo "$version"
}

# Compares 2 python versions in major, minor and micro parts
# $1 - version #1
# $2 - version #2
# Returns:
# 0 - version #1 is newer
# 1 - version #2 is newer
# 2 - versions are equal
version_compare() {
    version_1="$1"
    version_2="$2"
    # ZSH Only:
    version_1=("${(@s:.:)version_1}")
    version_2=("${(@s:.:)version_2}")
    major_1=$version_1[1]
    major_2=$version_2[1]
    minor_1=$version_1[2]
    minor_2=$version_2[2]
    micro_1=$version_1[3]
    micro_2=$version_2[3]
    # POSIX, but slow:
    # major_1="$(echo "$version_1" | cut -d. -f1)"
    # major_2="$(echo "$version_2" | cut -d. -f1)"
    # minor_1="$(echo "$version_1" | cut -d. -f2)"
    # minor_2="$(echo "$version_2" | cut -d. -f2)"
    # micro_1="$(echo "$version_1" | cut -d. -f3)"
    # micro_2="$(echo "$version_2" | cut -d. -f3)"

    # Compare majors
    if [ $major_1 -gt $major_2 ]; then
        # version 1's major is bigger, version 1 is newer
        return 0
    elif [ $major_1 -lt $major_2 ]; then
        # version 1's major is smaller, version 2 is newer
        return 1
    fi

    # Majors equal, compare minors
    if [ $minor_1 -gt $minor_2 ]; then
        # version 1's minor is bigger, version 1 is newer
        return 0
    elif [ $minor_1 -lt $minor_2 ]; then
        # version 1's major is smaller, version 2 is newer
        return 1
    fi

    # Minors equal, compare micros
    if [ $micro_1 -gt $micro_2 ]; then
        # version 1's micro is bigger, version 1 is newer
        return 0
    elif [ $micro_1 -lt $micro_2 ]; then
        # version 1's micro is smaller, version 2 is newer
        return 1
    fi

    # Micros equal, versions equal
    return 2
}

# Define new aliases if they don't already exsist, in which case override
# if the current version is newer than the version in the alias
# $1 - full valid pyenv python version (for example '3.6.12', `3.11-dev`, or `pypy3.6-7.2.0-src`)
# $2 - version used in the alias (for example '3.6', '3', or even '', but also `pypy3.7`, ...)
try_define_aliases() {
    version="$1"
    alias_version="$2"

    # Check if alias already exists
    if alias "py$alias_version" >/dev/null; then
        # Compare version from the existing alias with current version,
        # if current is newer, override the existing alias(es)
        defined_version="$(get_alias_version "$alias_version")"
        if version_compare "$version" "$defined_version"; then
            define_aliases "$version" "$alias_version"
            # echo "Overwrote '$alias_version' aliases to point to '$version'"
            return 0
        else
            return 1
        fi
    fi

    # The aliases aren't already defined, it's safe to create them
    define_aliases "$version" "$alias_version"
    # echo "Made '$alias_version' aliases pointing to '$version'"
    return 0
}

define_version_aliases() {
    prefix="$1"
    version="$2"
    suffix="$3"


    # ZSH only:
    version_data=("${(@s:.:)version}")
    major_version=$version_data[1]
    minor_version=$version_data[2]
    # POSIX, but slow:
    # major_version="$(echo "$version" | cut -d. -f1)"
    # minor_version="$(echo "$version" | cut -d. -f2)"

    # Define the major.minor.micro (full) alias
    try_define_aliases "$version" "$prefix$version$suffix"
    # Define the major.minor alias
    try_define_aliases "$version" "$prefix$major_version.$minor_version$suffix"
    # Define the major alias
    try_define_aliases "$version" "$prefix$major_version$suffix"
    # Define top level alias
    try_define_aliases "$version" "$prefix$suffix"
}

for python_dir in "$PYENV_ROOT"/versions/*/ ; do
    full_version="$(basename $python_dir)"

    version_info="$(parse_python_version "$full_version")"
    if [ $version_info = ';;' ]; then
        # Version info wasn't obtained successfully, skip this version
        echo "Skipped $full_version"
        continue
    fi

    # ZSH only:
    version_data=("${(@s:;:)version_info}")
    prefix=$version_data[1]
    version=$version_data[2]
    suffix=$version_data[3]
    # POSIX, but slow:
    # prefix="$(echo "$version_info" | cut -d';' -f1)"
    # version="$(echo "$version_info" | cut -d';' -f2)"
    # suffix="$(echo "$version_info" | cut -d';' -f3)"

    # startTime=$(date +%N)
    define_version_aliases "$prefix" "$version" "$suffix"
    # endTime=$(date +%N)
    # nanos="$(expr $endTime - $startTime)"
    # echo "took $(expr $nanos / 1000000) miliseconds"
done

if command -v poetry >/dev/null 2>&1; then
    alias poetry-pyenv='poetry env use "$(pyenv which python)" && poetry install'
fi