#!/usr/bin/env bash
unset CDPATH

if [ -d /C ]; then # Windows
home="$USERPROFILE" # NOTE: checked for existence later instead of here.
downloads="$home\\Downloads"
base="$home\\mamba"
else # Linux
home="$HOME" # NOTE: checked for existence later instead of here.
downloads="$home/src"
base="$home/mamba"
fi

# micromamba stuff; doesn't affect python itself:
installer_version="1.5.7" # remember to update both sha256s if you change this!
installer_archive="micromamba-$installer_version-0.tar.bz2"

releases="https://github.com/mamba-org/micromamba-releases/releases/download"
if [ -d /C ]; then # Windows
installer_sha256="c62a6a273e9a2b0313d2d6e50f27222447b3a2005dc4b9fef83ccd50b5ccce64"
installer_remote="$releases/$installer_version-0/micromamba-win-64.tar.bz2"
installer_path="Library/bin/micromamba.exe"
else # Linux
installer_sha256="2b92516aa575b42467805f25315bbd01454b8c7ec59ce04a5092187c4958c9a2"
installer_remote="$releases/$installer_version-0/micromamba-linux-64.tar.bz2"
installer_path="bin/micromamba"
fi

# remember not to use == when you mean = here:
python='python=3.11'
channel=conda-forge

main_packages=(
    'pip>=23.3.1' # python package manager
)

# pip packages should be pinned with pure equality
# because they're not exported by mamba env export.
pypi_packages=(
    'hipsterplot==0.1' # terminal plotting
)

# end basic settings, begin functions:

check_installer() {
    [ -s "$1" ] || return
    printf '%s %s\n' "$installer_sha256" "$1" | sha256sum -c - || return
}

mamba_paths() {
    if [ -d /C ]; then # Windows
        # TODO: infer System32 path by envvars. (it's looking for chcp.com)
        printf %s \
        "$(cygpath -u "C:\\Windows\\System32")" \
        : \
        "$(cygpath -u "C:\\Windows\\")" \
        ;
    else # Linux
        printf %s \
        "/usr/local/sbin" \
        : \
        "/usr/local/bin" \
        : \
        "/usr/sbin" \
        : \
        "/usr/bin" \
        : \
        "/sbin" \
        : \
        "/bin" \
        ;
    fi
}

python_paths() { # relies on globals: fullbase, env
    if [ -d /C ]; then # Windows
        printf %s \
        "$(cygpath -u "$fullbase\\envs\\$env")" \
        : \
        "$(cygpath -u "$fullbase\\envs\\$env\\Library\\mingw-w64\\bin")" \
        : \
        "$(cygpath -u "$fullbase\\envs\\$env\\Library\\bin")" \
        ;
    else # Linux
        printf %s \
        "$(readlink -f "$fullbase/envs/$env/x86_64-conda-linux-gnu/bin")" \
        : \
        "$(readlink -f "$fullbase/envs/$env/bin")" \
        ;
    fi
    printf %s "${MOREPATH:+:}$MOREPATH"
}

run_mamba() { # relies on globals: fullbase, base, env
    local paths=
    paths="$(mamba_paths)" || return
    PATH="$paths" "$fullbase/$installer_path" --no-rc --no-env -r "$base" "$@"
}

run_python() { # relies on globals: fullbase, env
    local exe= paths=
    [ -d /C ] && exe=python.exe || exe=python
    paths="$(python_paths)" || return
    PATH="$paths" "$exe" "$@"
}

find_msvc() { # TODO: does this use any globals? besides SYSTEMDRIVE which is an envvar.
    [ -d /C ] || { echo 'Internal error: attempted Windows routine on Linux!' >&2; exit 1; }
    [ -n "$SYSTEMDRIVE" ] || { echo '$SYSTEMDRIVE must be set.' >&2; exit 1; }

    dirs= libs= lib= includes=
    local dir= temp=
    local profiles="Program Files (x86)"

    cd "$SYSTEMDRIVE" || return
    [ -d "$profiles" ] || return

    add_dir() {
        dir="$1"
        dir="${dir%/}"
        dir="$(readlink -f "$dir")" || { printf 'readlink failed: %s\n' "$1" >&2; return 1; }
        if [ ! -d "$dir" ] && [ -e "$dir" ]; then
            dir="${dir%/*}"
        fi
        [ -d "$dir" ] || { printf 'not a directory: %s\n' "$dir"; return 1; }
        dirs+=":$dir"
    }

    windir() {
        # TODO: don't use cygpath, somehow.
        temp="$(cygpath -w "$1")" || return
        [ -d "$temp" ] || return
    }

    dir="$profiles/Microsoft Visual Studio"
    [ -d "$dir" ] || return
    dir="$(find "$dir" -name cl.exe | grep -E '[\/]Hostx64[\/]x64[\/]cl.exe$' | LC_ALL=C sort | tail -n1)"
    add_dir "$dir" || return

    windir "${dir%/bin/*}/lib/x64" || return; libs+=";$temp"; lib="$temp"
    windir "${dir%/bin/*}/include" || return; includes+=";$temp"

    dir="$profiles/Windows Kits"
    [ -d "$dir" ] || return
    dir="$(find "$dir" -name rc.exe | grep -E '[\/]x64[\/]rc.exe' | LC_ALL=C sort | tail -n1)"
    add_dir "$dir" || return

    # NOTE: this next line relies on add_dir behavior a little.
    windir "$dir/ucrt" || return; includes+=";$temp"

    dir="${dir%/x64}"
    windir "${dir/bin/Lib}/um/x64" || return; libs+=";$temp"
    windir "${dir/bin/Lib}/ucrt/x64" || return; libs+=";$temp"
    windir "${dir/bin/Include}/ucrt" || return; includes+=";$temp"
    windir "${dir/bin/Include}/shared" || return; includes+=";$temp"

    dirs="${dirs#:}"
    libs="${libs#;}"
    includes="${includes#;}"
    #printf '%s\n' "DIRS: $dirs" >&2
    #printf '%s\n' "LIBS: $libs" >&2
    #printf '%s\n' "INCLUDES: $includes" >&2
    #exit 1
}

install_msvc() { # TODO: document globals used.
    [ -d /C ] || { echo 'Internal error: attempted Windows routine on Linux!' >&2; exit 1; }
    [ -n "$SYSTEMDRIVE" ] || { echo '$SYSTEMDRIVE must be set.' >&2; exit 1; }
    [ -n "$COMSPEC" ] || { echo '$COMSPEC must be set.' >&2; exit 1; }

    # find the latest version: curl -I 'https://aka.ms/vs/17/release/vs_buildtools.exe'
    local uuid='5bebe58c-9308-4a5b-9696-b6f84e90a32e'
    local hash='d62702bf9e2bb2c8be1f85ec4b86e0426e42646d12ac5196c451574d22be148e'
    local exe='vs_BuildTools.exe'
    local installer="$downloads\\$exe"
    local url="https://download.visualstudio.microsoft.com/download/pr/$uuid/$hash/$exe"
    local profiles="Program Files (x86)"

    mode=install
    if [ "$1" = update ]; then
        shift
        mode=update
    fi

    # always download the installer since it's only a couple megabytes.
    curl -L "$url" -o "$installer" || return
    printf '%s %s\n' "$hash" "$installer" | sha256sum -c - || return

    cd "$SYSTEMDRIVE" || return
    [ -d "$profiles" ] || return

    local bt
    bt="$profiles\\Microsoft Visual Studio\\2022\\BuildTools"

    # TODO: use "update" instead of "install" when applicable.
    # despite the name, you can use the Windows 11 SDK (or at least this version of it)
    # to target as early as Windows 10 v1507.
    "$installer" --quiet --wait --norestart --nocache "$mode" \
        --add 'Microsoft.VisualStudio.Component.VC.Tools.x86.x64' \
        --add 'Microsoft.VisualStudio.Component.Windows11SDK.22000'
    local ret=$?

    # bash has run ERRORLEVEL through a modulus of 256, so some hacks are required.
    local code=0
    case $ret in
        (0)     code=0;; # success (probably)
        (58)    code=-1073741510;;
        (65)    code=8001;;
        (66)    code=1602;; # NOTE: this could also be 8002!
        (67)    code=8003;;
        (68)    code=8004;;
        (69)    code=8005;;
        (70)    code=8006;;
        (82)    code=1618;;
        (87)    code="INVALID ARGUMENTS ($ret)";;
        (105)   code=1641;;
        (139)   code=5003;;
        (140)   code=5004;;
        (141)   code=5005;;
        (143)   code=5007;;
        (145)   code=-1073720687;;
        (194)   code=3010;; # success; restart required (but not really)
        (228)   code=740;;
        (233)   code=1001;;
        (235)   code=1003;;
        (*)     code="UNKNOWN ($ret)";;
    esac

    printf '%s exited with %%ERRORLEVEL%% %s.\n' "$installer" "$code" >&2
    [ "$code" = 0 ] || [ "$code" = 3010 ] || return 1
}

run_with_compiler() { # relies on globals: fullbase, env, working_dir
    if [ -d /C ]; then # Windows
        find_msvc || install_msvc \
         || { printf 'failed to install MSVC and a Windows SDK.\n' >&2; exit 1; }
        find_msvc \
         || { printf 'failed to find a valid installation of MSVC.\n' >&2; exit 1; }
        # globals should now exist: dirs, libs, lib, includes

        (
            cd "$working_dir" || return
            export CMAKE_GENERATOR="Visual Studio 17 2022"
            MOREPATH="$dirs" LIB="$libs" LIBPATH="$lib" INCLUDE="$includes" \
            run_python "$@"
        )

    else # Linux
        printf '%s\n' 'run_with_compiler: TODO: handle static gcc here!' >&2
        run_python "$@"
    fi
}

run_pip() {
    run_with_compiler -m pip --disable-pip-version-check "$@"
}

find_last_env() { # relies on globals: fullbase
    local arg="$1" e= oe=
    env="$(date -u '+%Y%m%d')" || exit
    n=1
    while [ -d "$fullbase/envs/$env-$n" ]; do
        let n++
    done
    new_env="$env-$n"
    let n--
    if [ $n = 0 ]; then
        env=_
        if [ -d "$fullbase/envs" ]; then
            for e in "$fullbase/envs/"*; do
                oe="${e##*/}"
                #printf 'considering %s... compared to %s\n' "$oe" "$env" >&2
                if [ -d "$e" ] && [ "${oe#20}" != "$oe" ] && [[ "$oe" > "$env" ]]; then
                    #printf '%s is more recent!\n' "$oe"
                    env="$oe"
                fi
            done
        fi
        if [ "$env" = _ ]; then
            env=ERROR
            if [ "$arg" != new ]; then
                echo 'failed to find most recent environment' >&2
                exit 1
            fi
        fi
    else
        env="$env-$n"
    fi
}

# end of functions, begin main:

if [ -d /C ]; then # Windows
    [ -n "$USERPROFILE" ] || { echo '$USERPROFILE must be set.' >&2; exit 1; }
else # Linux
    [ -n "$HOME" ] || { echo '$HOME must be set.' >&2; exit 1; }
fi

[ -d "$downloads" ] || mkdir "$downloads" || exit

working_dir="$PWD"
[ -d "$working_dir" ] || { echo 'failed to determine working directory' >&2; exit 1; }

[ -d "$base" ] || mkdir "$base" || exit
if [ -d /C ]; then # Windows
    fullbase="$(cygpath -u "$base")" || exit
    fullbase="$(readlink -f "$fullbase")" || exit
    fullbase="$(cygpath -w "$fullbase")" || exit
else # Linux
    fullbase="$(readlink -f "$base")" || exit
fi

version="unknown"
if [ -s "$fullbase/$installer_path" ]; then
    # TODO: do something meaningful if the execution itself fails.
    version="$(run_mamba --version)"
fi

if [ "$version" != "$installer_version" ]; then
    archive="$downloads/$installer_archive"
    if check_installer "$archive"; then
        :
    else
        # slightly less safe: curl -LOJ "$installer_remote"
        curl -L "$installer_remote" -o "$archive" || exit
        check_installer "$archive" || exit
    fi
    ! [ -d /C ] || archive="$(cygpath -u "$archive")" || exit
    cd "$fullbase" || exit
    if ! which bzip2 >/dev/null 2>&1 && which python3 >/dev/null 2>&1; then
        # ubuntu-server does not come with bzip2, but python3 can save us here.
        # 'import tarfile,sys;_,a,b,c=sys.argv;tarfile.open(a,"r:bz2").extract(b,c)
        python3 -c 'import tarfile,sys;_,a,b=sys.argv;tarfile.open(a,"r:bz2").extract(b)' \
        "$archive" "$installer_path" || exit
    else
        tar jxf "$archive" "$installer_path" || exit
    fi
    cd "$working_dir" || exit
fi

[ "$1" = --create ] && creating=1 || creating=0
[ "$1" = --mamba ] && dancing=1 || dancing=0
[ "$1" = --pip ] && compiling=1 || compiling=0
[ "$1" = --wheel ] && wheeling=1 || wheeling=0
[ "$1" = --msvc ] && visualizing=1 || visualizing=0

if [ "${AM_ENV:-__empty__}" = __empty__ ]; then
    find_last_env new || exit
else
    if [ $creating = 1 ]; then
        :
    elif [ ! -d "$fullbase/envs/$AM_ENV" ]; then
        printf 'failed to find environment $AM_ENV: %s\n' "$AM_ENV" >&2
        sleep 3 #exit 1
    elif [ ! -s "$fullbase/envs/$AM_ENV/LICENSE_PYTHON.txt" ]; then
        printf 'empty environment $AM_ENV: %s\n' "$AM_ENV" >&2
        sleep 3 #exit 1
    fi
    env="$AM_ENV"
fi

if [ $visualizing = 1 ]; then
    [ $visualizing = 0 ] || shift || exit
    install_msvc update
    exit
fi

if [ $dancing = 1 ]; then
    [ $dancing = 0 ] || shift || exit

    no_default=0
    for arg in "$@"; do
        [ "$arg" != --name ] || no_default=1
        # TODO: is --name=stuff valid? detect that too.
        if [ "${arg#--}" = "$arg" ]; then # ignore other long arguments
            # detect -n, -asdfn shortargs. startswith -, endswith n.
            if [ "${arg#-}" != "$arg" ] && [ "${arg%n}" != "${arg}" ]; then
                no_default=1
            fi
        fi
    done
    if [ $no_default = 0 ]; then
        run_mamba -n "$env" "$@" || exit
    else
        run_mamba "$@" || exit
    fi

    [ $dancing = 0 ] || exit 0
fi

if [ $compiling = 1 ]; then
    [ $compiling = 0 ] || shift || exit

    [ "$env" != ERROR ] || { printf 'failed to find most recent environment\n' >&2; exit 1; }
    run_pip "$@" || exit

    [ $compiling = 0 ] || exit 0
fi

if [ $wheeling = 1 ]; then
    [ $wheeling = 0 ] || shift || exit

    [ "$env" != ERROR ] || { printf 'failed to find most recent environment\n' >&2; exit 1; }
    run_with_compiler -u setup.py bdist_wheel || exit

    [ $wheeling = 0 ] || exit 0
fi

if [ "$env" = ERROR ] || [ $creating != 0 ]; then
    [ $creating = 0 ] || shift || exit
    [ -z "$new_env" ] || env="$new_env"

    if [ $creating = 0 ]; then
        printf 'WARNING: automatically creating a new environment: %s\n' "$env" >&2
    fi

    _do_create() {
        if [ $# = 0 ]; then
            run_mamba -yqn "$env" create -c "$channel" "$python" "${main_packages[@]}"
        else
            run_mamba -yqn "$env" create "$@"
        fi
    }

    if [ $creating = 1 ]; then
        _do_create "$@" 2>&1 | tee "$fullbase/create.log"
    else
        _do_create 2>&1 | tee "$fullbase/create.log"
    fi
    res=${PIPESTATUS[0]}
    [ $res = 0 ] || exit "$res"

    if [ $# = 0 ]; then
        common=(--isolated --disable-pip-version-check)
        run_pip install "${common[@]}" --no-deps "${pypi_packages[@]}" 2>&1 \
          | tee -a "$fullbase/create.log"
        [ ${PIPESTATUS[0]} = 0 ] || exit ${PIPESTATUS[0]}
        # using run_python is faster than run_pip since it won't check for compilers.
        run_python -m pip "${common[@]}" check \
          | tee -a "$fullbase/create.log"
        #[ ${PIPESTATUS[0]} = 0 ] || exit ${PIPESTATUS[0]}
    fi

    pre="$fullbase/$env"

    run_mamba -n "$env" env export > "$pre.yml" || exit

    run_mamba -n "$env" list --json | run_python -c \
    'for(p)in(m:=__import__)("json").load(m("sys").stdin):print(p["name"]+"=="+p["version"])' \
      > "$pre.conda.freeze.txt" || exit
    [ ${PIPESTATUS[0]} = 0 ] || exit ${PIPESTATUS[0]}

    run_python -m pip "${common[@]}" list --pre --format=freeze \
      > "$pre.pip.freeze.txt" || exit

    echo '- pip:' | cat "$pre.yml" - > "$pre.with-pip.yml" || exit
    run_python -c \
    '[print(*" -",b[k])for(a,b)in[map(lambda f:{l.strip().lower().replace(*"-_"):l.strip()for(l)in open(f)},__import__("sys").argv[1:])]for(k)in(b)if(k)not in a]' \
    "$pre.conda.freeze.txt" "$pre.pip.freeze.txt" \
     >> "$pre.with-pip.yml" || exit

    rm "$pre.conda.freeze.txt" "$pre.pip.freeze.txt" || exit

    [ $creating = 0 ] || exit 0
fi

run_python "$@"