#!/bin/bash

# PEnv, dead simple python venv manager. Handles creating the env and installing the jupyter spec.
# Version 0.1 - 20th Nov, 2025
# MIT Licensed, by The Grammer Society

# Issues:
# names are fully lower case when installed as jupyterhub spec. convert name to lower case when deleting an env.
# when uninstalling, remove any jupyter specs. iterate through envs in .env, call jspec remove name.

usage() {
    echo ""
    echo "Usage: penv <command> <args>"
    echo "Commands:"
    echo " - new <name> <ver=default> - Create a new environment, optional, defaults to system's default python"
    echo "                              ex: penv new tools 3.13 -> creates new environment by the name tools, using Python 3.13"
    echo "                              ex: penv new main       -> creates new env with the name main, using the default Python version on the system"
    echo " - del <name>               - Delete existing environment"
    # echo " - setdef <name>            - Use an existing environment as the default kernel"
    echo " - activate <name>          - Activates an existing environment (depricated, use pact)"
    echo " - list                     - List available environments"
    echo " - pylist                   - List available PYTHON versions"
    echo " - install                  - Initialize user's profile"
    echo " - uninstall                - Remove ALL the environments installed using PENV"
    echo " - help                     - Print this message"
    echo
    echo "Aliases:"
    echo " - new      -> mk"
    echo " - del      -> rm"
    # echo " - setdef   -> sd"
    echo " - activate -> ac, act"
    echo " - pylist   -> pyls"
    echo " - list     -> ls"
}

title(){
    echo "PEnv: Python Environment Manager for Local Jupyter Instances"
}

create(){
    name=$1
    vers=$2
    if [[ -z "$name" ]]; then
        echo Provide a name for the environment
        exit 2
    fi
    if [[ -d "$HOME/.envs/$name" ]]; then
        echo Environment \""$name"\" already exists.
        exit 4
    fi

    if [[ -z "$vers" ]]; then
        if [ ! -f /usr/bin/python3 ]; then
            echo "Base python3 not found on your system. Install it using your package manager"
            exit 3
        fi
        vers=$(/usr/bin/python3 -V | cut -d " " -f 2 | cut -d "." -f 1,2)
        echo "No version mentioned, using Python $vers"
    fi
    if [ ! -f /usr/bin/python"$vers" ]; then
        echo "Python $vers not installed on system."
        exit 7
    fi
    mkdir -p "$HOME"/.envs
    /usr/bin/python"$vers" -m venv "$HOME"/.envs/"$name"
    source "$HOME"/.envs/"$name"/bin/activate
    python -m pip install ipykernel ipywidgets
    python -m ipykernel install --user --name "$name"
    echo
    echo "--------------------------------------------------------------"
    echo "Created \"$name\". It should be available as a Jupyter Kernel."
    echo "Use \"pact $name\" to activate in command line."
    deactivate
}

delete(){
    name=$1
    if [ ! -d "$HOME/.envs/$name" ]; then
        echo "Environment \"$name\" not found"
        exit 5
    fi
    if [[ "$VIRTUAL_ENV" == "$HOME/.envs/$name" ]]; then
        echo "Env still active, deactivate it before removing it."
        echo "Run \"deactivate\" to deactivate it"
        exit 6
    fi
    read -r -n 1 -p "Remove environment $name (y/N)? " choise
    case $choise in
    [yY])
        size=$(du "$HOME"/.envs/"$name" -sh | cut -f1)
        source "$HOME"/.envs/"$name"/bin/activate
        echo
        echo Spec $(jupyter-kernelspec remove -f -y "$name")
        deactivate
        rm -r "$HOME/.envs/$name"
        echo "Data Removed $HOME/.envs/$name"
        echo "Reclaimed $size"
        ;;
    [nN])
        echo ""
        echo "Cancelled"
        ;;
    *)
        echo ""
        echo "Invalid input"
        ;;
    esac
}

# setdef() {
#     name=$1
#     if [ ! -d "$HOME/.envs/$name" ]; then
#         echo "Env $name doesn't exist."
#         exit 9
#     fi
#     source $HOME/.envs/$name/bin/activate
#     python -m ipykernel install --user
#     deactivate
# }

activate() {
    echo Use \"pact "$1"\"
}

list(){
    if [ ! -d "$HOME/.envs" ]; then
        echo "PENV not initialized. Run \"penv init\""
        exit 8
    fi
    ls "$HOME"/.envs
}

install() {
    if [[ $(realpath "$0") == "$HOME/.local/bin/penv" ]]; then
        echo PENV already installed
        exit 10
    fi
    mkdir -p "$HOME/.local/bin"
    cp $(realpath "$0") "$HOME"/.local/bin/
    mkdir -p "$HOME"/.envs/
    echo "
    pact () {
        if [[ -z \"\$1\" ]]; then
            echo \"Usage: pact \<name\>\"
            return;
        fi;
        if [ ! -d \"\$HOME/.envs/\$1\" ]; then
            echo \"Env \$1 not found\"
            return;
        fi
        . \$HOME/.envs/\$1/bin/activate
    }
    " > "$HOME"/.envs/.penv.funcs
    add2RC bashrc
    add2RC zshrc
    echo "Initialized PENV. Restart the shell"
}

add2RC(){
    if [[ -f "$HOME/.$1" ]]; then
        grep -qxF "source \$HOME/.envs/.penv.funcs" "$HOME/.$1" || echo "source \$HOME/.envs/.penv.funcs" >> "$HOME/.$1"
    fi
}

uninstall() {
    if [[ -n "$VIRTUAL_ENV" ]]; then
        echo "Deactivate the existing environment before uninstalling"
        exit 11
    fi
    read -r -n 1 -p "REALLY UNINSTALL ALL THE ENVIRONMENTS?(y/N) " ans
    case $ans in
        [yY])
            echo
            ;;
        [nN])
            echo
            echo "Cancelled"
            return;
            ;;
        *)
            echo
            echo "Invalid response"
            return;
            ;;
    esac
    sed -i "/source \$HOME\/.envs\/.penv.funcs/d" "$HOME"/.bashrc
    rm -r "$HOME"/.envs
    rm "$HOME"/.local/bin/penv
    echo "Uninstalled PENV"
}

pyls(){
    find /usr/local/bin /usr/bin -name "python3.*" 2> /dev/null | tr " " "\n" | grep "3.*" --color=always
}

if [ "$#" -lt 1 ]; then
    title
    usage
    exit 1
fi

cmd=$1

case "$cmd" in
    new|mk)
        create "$(echo "$2" | tr "[:upper:]" "[:lower:]")" "$3"
        ;;
    del|rm)
        delete "$(echo "$2" | tr "[:upper:]" "[:lower:]")"
        ;;
    # setdef|sd)
    #     setdef $2
    #     ;;
    act|ac)
        activate "$(echo "$2" | tr "[:upper:]" "[:lower:]")"
        ;;
    list|ls)
        list
        ;;
    pylist|pyls)
        pyls
        ;;
    help)
        title
        usage
        ;;
    install)
        install
        ;;
    uninstall)
        uninstall
        ;;
    *)
        echo "invalid command: \"$1\"" >&2
        echo "$SCRIPT"
        usage
        ;;
esac