set fallback := true

tempdir := `mktemp -d`
export JUPYTERHUB_NAMESPACE := env("JUPYTERHUB_NAMESPACE", "jupyter")
export JUPYTERHUB_CHART_VERSION := env("JUPYTERHUB_CHART_VERSION", "4.2.0")
export JUPYTERHUB_OIDC_CLIENT_ID := env("JUPYTERHUB_OIDC_CLIENT_ID", "jupyterhub")
export JUPYTERHUB_ENABLE_NFS_PV := env("JUPYTERHUB_ENABLE_NFS_PV", "")
export JUPYTERHUB_VAULT_INTEGRATION_ENABLED := env("JUPYTERHUB_VAULT_INTEGRATION_ENABLED", "false")
export JUPYTER_PYTHON_KERNEL_TAG := env("JUPYTER_PYTHON_KERNEL_TAG", "python-3.12-4")
export KERNEL_IMAGE_BUUN_STACK_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_REPOSITORY", "buun-stack-notebook")
export KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY := env("KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY", "buun-stack-cuda-notebook")
export JUPYTER_PROFILE_MINIMAL_ENABLED := env("JUPYTER_PROFILE_MINIMAL_ENABLED", "false")
export JUPYTER_PROFILE_BASE_ENABLED := env("JUPYTER_PROFILE_BASE_ENABLED", "false")
export JUPYTER_PROFILE_DATASCIENCE_ENABLED := env("JUPYTER_PROFILE_DATASCIENCE_ENABLED", "true")
export JUPYTER_PROFILE_PYSPARK_ENABLED := env("JUPYTER_PROFILE_PYSPARK_ENABLED", "false")
export JUPYTER_PROFILE_PYTORCH_ENABLED := env("JUPYTER_PROFILE_PYTORCH_ENABLED", "false")
export JUPYTER_PROFILE_TENSORFLOW_ENABLED := env("JUPYTER_PROFILE_TENSORFLOW_ENABLED", "false")
export JUPYTER_PROFILE_BUUN_STACK_ENABLED := env("JUPYTER_PROFILE_BUUN_STACK_ENABLED", "false")
export JUPYTER_PROFILE_BUUN_STACK_CUDA_ENABLED := env("JUPYTER_PROFILE_BUUN_STACK_CUDA_ENABLED", "false")
export IMAGE_REGISTRY := env("IMAGE_REGISTRY", "localhost:30500")
export KEYCLOAK_REALM := env("KEYCLOAK_REALM", "buunstack")
export LONGHORN_NAMESPACE := env("LONGHORN_NAMESPACE", "longhorn")
export VAULT_ADDR := env("VAULT_ADDR", "http://vault.vault.svc:8200")

[private]
default:
    @just --list --unsorted --list-submodules

# Add Helm repository
add-helm-repo:
    helm repo add jupyterhub https://jupyterhub.github.io/helm-chart
    helm repo update

# Remove Helm repository
remove-helm-repo:
    helm repo remove jupyterhub

# Create JupyterHub namespace
create-namespace:
    kubectl get namespace ${JUPYTERHUB_NAMESPACE} &>/dev/null || \
        kubectl create namespace ${JUPYTERHUB_NAMESPACE}

# Delete JupyterHub namespace
delete-namespace:
    kubectl delete namespace ${JUPYTERHUB_NAMESPACE} --ignore-not-found

# Install JupyterHub
install:
    #!/bin/bash
    set -euo pipefail
    export JUPYTERHUB_HOST=${JUPYTERHUB_HOST:-}
    while [ -z "${JUPYTERHUB_HOST}" ]; do
        JUPYTERHUB_HOST=$(
            gum input --prompt="JupyterHub host (FQDN): " --width=100 \
            --placeholder="e.g., jupyter.example.com"
        )
    done
    just create-namespace
    # just k8s::copy-regcred ${JUPYTERHUB_NAMESPACE}
    just keycloak::create-client ${KEYCLOAK_REALM} ${JUPYTERHUB_OIDC_CLIENT_ID} \
        "https://${JUPYTERHUB_HOST}/hub/oauth_callback"
    just add-helm-repo
    export JUPYTERHUB_OIDC_CLIENT_ID=${JUPYTERHUB_OIDC_CLIENT_ID}
    export KEYCLOAK_REALM=${KEYCLOAK_REALM}
    export JUPYTER_PYTHON_KERNEL_TAG=${JUPYTER_PYTHON_KERNEL_TAG}
    export JUPYTER_FSGID=${JUPYTER_FSGID:-100}
    export PVC_NAME=""
    if [ -z "${JUPYTERHUB_ENABLE_NFS_PV}" ]; then
        if gum confirm "Are you going to use NFS PV?"; then
            JUPYTERHUB_ENABLE_NFS_PV=true
        else
            JUPYTERHUB_ENABLE_NFS_PV=false
        fi
    fi
    if [ "${JUPYTERHUB_ENABLE_NFS_PV}" = "true" ]; then
        if ! helm status longhorn -n ${LONGHORN_NAMESPACE} &>/dev/null; then
            echo "Longhorn is not installed. Please install Longhorn first." >&2
            exit 1
        fi
        export JUPYTER_NFS_IP=${JUPYTER_NFS_IP:-}
        while [ -z "${JUPYTER_NFS_IP}" ]; do
            JUPYTER_NFS_IP=$(
                gum input --prompt="NFS server IP address: " --width=100 \
                --placeholder="e.g., 192.168.10.1"
            )
        done
        export JUPYTER_NFS_PATH=${JUPYTER_NFS_PATH:-}
        while [ -z "${JUPYTER_NFS_PATH}" ]; do
            JUPYTER_NFS_PATH=$(
                gum input --prompt="NFS server export path: " --width=100 \
                --placeholder="e.g., /volume1/drive1/jupyter"
            )
        done
        PVC_NAME=jupyter-nfs-pvc
        if ! kubectl get pv jupyter-nfs-pv &>/dev/null; then
            gomplate -f nfs-pv.gomplate.yaml | kubectl apply -f -
        fi
        kubectl apply -n ${JUPYTERHUB_NAMESPACE} -f nfs-pvc.yaml
    fi
    # https://z2jh.jupyter.org/en/stable/
    gomplate -f jupyterhub-values.gomplate.yaml -o jupyterhub-values.yaml
    helm upgrade --cleanup-on-fail --install jupyterhub jupyterhub/jupyterhub \
        --version ${JUPYTERHUB_CHART_VERSION} -n ${JUPYTERHUB_NAMESPACE} \
        --timeout=20m -f jupyterhub-values.yaml
    # wait deployments manually because `helm upgrade --wait` does not work for JupyterHub
    just k8s::wait-deployments-ready ${JUPYTERHUB_NAMESPACE} hub proxy

    # Setup Vault integration if enabled
    if [ "${JUPYTERHUB_VAULT_INTEGRATION_ENABLED}" = "true" ]; then
        just setup-vault-jwt-auth
    fi

# Uninstall JupyterHub
uninstall:
    #!/bin/bash
    set -euo pipefail
    helm uninstall jupyterhub -n ${JUPYTERHUB_NAMESPACE} --wait --ignore-not-found
    kubectl delete pods -n ${JUPYTERHUB_NAMESPACE} -l app.kubernetes.io/component=singleuser-server
    kubectl delete -n ${JUPYTERHUB_NAMESPACE} pvc jupyter-nfs-pvc --ignore-not-found
    if kubectl get pv jupyter-nfs-pv &>/dev/null; then
        kubectl patch pv jupyter-nfs-pv -p '{"spec":{"claimRef":null}}'
    fi

# Delete JupyterHub PV
delete-pv:
    #!/bin/bash
    set -euo pipefail
    if kubectl get pv jupyter-nfs-pv &>/dev/null; then
        kubectl patch pv jupyter-nfs-pv -p '{"spec":{"claimRef":null}}'
        kubectl delete pv jupyter-nfs-pv
    fi

# Build Jupyter notebook kernel images
build-kernel-images:
    #!/bin/bash
    set -euo pipefail
    # Build python package wheel
    cd ../python-package
    rm -rf dist/ build/ *.egg-info/
    SETUPTOOLS_SCM_PRETEND_VERSION_FOR_BUUNSTACK=0.1.0 python -m build --wheel
    cd ../jupyterhub
    # Copy built wheel to image directories
    cp ../python-package/dist/*.whl ./images/datastack-notebook/
    cp ../python-package/dist/*.whl ./images/datastack-cuda-notebook/
    (
        cd ./images/datastack-notebook
        docker build -t \
        ${IMAGE_REGISTRY}/${KERNEL_IMAGE_BUUN_STACK_REPOSITORY}:${JUPYTER_PYTHON_KERNEL_TAG} \
        --build-arg spark_version="3.5.4" \
        --build-arg spark_download_url="https://archive.apache.org/dist/spark/" \
        .
    )
    (
        cd ./images/datastack-cuda-notebook
        docker build -t \
        ${IMAGE_REGISTRY}/${KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY}:${JUPYTER_PYTHON_KERNEL_TAG} \
        --build-arg spark_version="3.5.4" \
        --build-arg spark_download_url="https://archive.apache.org/dist/spark/" \
        .
    )
    # Clean up copied wheel files
    rm -f ./images/datastack-notebook/*.whl
    rm -f ./images/datastack-cuda-notebook/*.whl

# Push Jupyter notebook kernel images
push-kernel-images: build-kernel-images
    docker push ${IMAGE_REGISTRY}/${KERNEL_IMAGE_BUUN_STACK_REPOSITORY}:${JUPYTER_PYTHON_KERNEL_TAG}
    docker push ${IMAGE_REGISTRY}/${KERNEL_IMAGE_BUUN_STACK_CUDA_REPOSITORY}:${JUPYTER_PYTHON_KERNEL_TAG}

# Configure Vault for JupyterHub integration
setup-vault-integration:
    #!/bin/bash
    set -euo pipefail
    echo "Creating JupyterHub Vault policy..."
    just vault::write-policy jupyter-user $(pwd)/vault-policy.hcl
    echo "✓ JupyterHub policy created"

# Setup JWT auth for JupyterHub tokens (no re-authentication needed)
setup-vault-jwt-auth:
    #!/bin/bash
    set -euo pipefail
    echo "Setting up Vault integration for JupyterHub..."
    just setup-vault-integration
    just vault::setup-jwt-auth "jupyterhub" "jupyter-token" "jupyter-user"
    echo "✓ Vault integration configured"
    echo ""
    echo "Users can now access Vault from notebooks using:"
    echo "  import os, hvac"
    echo "  client = hvac.Client(url=os.getenv('VAULT_ADDR'), verify=False)"
    echo "  client.auth.jwt.jwt_login("
    echo "      role='jupyter-token',"
    echo "      jwt=os.getenv('JUPYTERHUB_OIDC_ACCESS_TOKEN'),"
    echo "      path='jwt'"
    echo "  )"