#!/bin/bash
#
# Start a pod session with Jupyter-notebook and port-forwarding.
# For KPU workshop only.
#
# Copyright: Charles Deledalle, 2019.
#

colorize () {
    a=$(echo -e "\x1b[29;01m")
    b=$(echo -e "\x1b[32;01m")
    c=$(echo -e "\e[m")
    sed "s/\([^#][^#]*\)/$a\1$c/g" | sed "s/\(##*\)/$b\1$c/g"
}

if [ "$1" == "--help" ] ; then
    echo -e "Usage: $0"
    sed '1,/^#$/d;/^#$/,$d;s/^# //' $0
    echo
    sed '1,/^#$/d;1,/^#$/d;/^#$/,$d;s/^# *//' $0
    exit 0
fi


#############################################
# Perform some verifications and setting up #
#############################################

# Check for dependencies
if ! hash kubectl 2> /dev/null ; then
    echo "Install kubectl"
    exit 2
fi
if ! hash netstat 2> /dev/null ; then
    echo "Install net-tools package"
    exit 2
fi
if ! hash nmap 2> /dev/null ; then
    echo "Install nmap package"
    exit 2
fi

# Prevent form Ctrl-C
trap "printf '\nPlease wait (or do \"kill -9 $$\")\n'" SIGINT SIGTERM

# Get kubernete user id and check if id is valid
user="$(kubectl config view --template='{{ range .contexts }}{{ if eq .name "'$(kubectl config current-context)'" }}{{ .context.user }}{{ end }}{{ end }}')"
user="`echo "$user" | sed 's!.*/!!'`"
if ! [ "$user" = "`echo "$user" | sed 's/[^0-9]//g'`" ] ; then
    echo Authentification failed
    exit 2
fi

# Get kubernete server hostname and port and check port is open for listening
server="$(kubectl config view --template='{{ range .clusters }}{{ if eq .name "'$(kubectl config current-context)'" }}{{ .cluster.server }}{{ end }}{{ end }}')"
hostname="`echo "$server" | sed 's!https*://\([^/]*\).*!\1!'`"
hostport="443"
if echo "$hostname" | grep -q ':[0-9][0-9]*$' ; then
    hostport="`echo "$hostname" | sed 's/.*:\([0-9][0-9]*\)$/\1/'`"
    hostname="`echo "$hostname" | sed 's/\(.*\):[0-9][0-9]*$/\1/'`"
fi
if ! nmap "$hostname" -p "$hostport" | grep "$hostport" | grep -q "open" ; then
    echo "Server not reachable (are you using a firewall?)"
    exit 2
fi
local_port=8888
while netstat --listening | grep localhost:$local_port -q ; do
    local_port=$(($local_port + 1))
done

# Setup variables
context=nautilus
namespace=kpuworkshop
podname="pod-u$user"
volname="vol-u$user"


##################
# Create context #
##################
mlineno=$LINENO
if ! kubectl config get-contexts --namespace=$namespace | grep $context -q ; then
    printf "\n\n"
    sed "$((mlineno-3)),$((mlineno-1))!d" $0 | colorize

    kubectl config set-context $context --namespace=$namespace

    sleep 5
else
    printf "Found existing context '$context' in '$namespace'.\n"
fi


###########################################
# Create persistent volume (about 5 secs) #
###########################################
mlineno=$LINENO
if ! kubectl get persistentvolumeclaims --namespace=$namespace 2> /dev/null | grep "$volname" -q ; then
    printf "\n\n"
    sed "$((mlineno-3)),$((mlineno-1))!d" $0 | colorize

    cat <<EOF | kubectl create -f -
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: $volname
spec:
  storageClassName: rook-block
  accessModes:
  - ReadWriteOnce
  resources:
    requests:
      storage: 5Gi
EOF
    sleep 5
    while ! kubectl get persistentvolumeclaims --namespace=$namespace 2> /dev/null | grep "$volname" -q ; do
	printf "Wait for volume...\n"
	sleep 1
    done

    # Create change permission pod
    if ! kubectl get pods --namespace=$namespace 2> /dev/null | grep "$podname-tmp" -q ; then
	cat <<EOF | kubectl create --namespace=$namespace -f -
apiVersion: v1
kind: Pod
metadata:
  name: $podname-tmp
spec:
  securityContext:
    runAsUser: 0
    runAsGroup: 0
    fsGroup: 0
  containers:
  - name: gpu-container
    image: gitlab-registry.nautilus.optiputer.net/prp/jupyterlab:latest
    command: ["sleep", "1h"]
    volumeMounts:
    - mountPath: /home/jovyan
      name: $volname
  restartPolicy: Never
  volumes:
  - name: $volname
    persistentVolumeClaim:
      claimName: $volname
EOF
	sleep 5
    else
	printf "Connecting to existing pod '$podname-tmp'.\n"
    fi
fi


#####################################
# Change permissions (about 5 secs) #
#####################################
mlineno=$LINENO
if kubectl get pods --namespace=$namespace 2> /dev/null | grep "$podname-tmp" -q ; then
    printf "\n\n"
    sed "$((mlineno-3)),$((mlineno-1))!d" $0 | colorize

    while ! kubectl get pods --namespace=$namespace --field-selector=status.phase=Running 2> /dev/null | grep "$podname-tmp" -q ; do
	printf "Wait for pod...\n"
	sleep 5
    done
    kubectl exec --namespace=$namespace -it "$podname-tmp" chown jovyan:users /home/jovyan
    kubectl delete pod --namespace=$namespace "$podname-tmp"
else
    printf "Using existing volume '$volname'.\n"
fi


####################################################
# Create pod (about 5 secs). Pod will finish in 5h #
####################################################
mlineno=$LINENO
if kubectl get pods --namespace=$namespace --field-selector=status.phase!=Running 2> /dev/null | grep "$podname" -q ; then
    kubectl delete pod "$podname"
fi
mountedpod="`kubectl describe persistentvolumeclaims "$volname" | grep "Mounted By:" | sed "s/Mounted By:[[:blank:]]*\(.*\)/\1/"`"
if [ "$mountedpod" != "<none>" ] && [ "$mountedpod" != "$podname" ] ; then
    kubectl delete pod "$mountedpod"
fi
if ! kubectl get pods --namespace=$namespace 2> /dev/null | grep "$podname" -q ; then
    printf "\n\n"
    sed "$((mlineno-3)),$((mlineno-1))!d" $0 | colorize

    cat <<EOF | kubectl create --namespace=$namespace -f -
apiVersion: v1
kind: Pod
metadata:
  name: $podname
spec:
  securityContext:
    runAsUser: 1000
    runAsGroup: 100
    fsGroup: 100
  containers:
  - name: gpu-container
    image: gitlab-registry.nautilus.optiputer.net/prp/jupyterlab:latest
    command: ["sleep", "5h"]
    resources:
      limits:
        nvidia.com/gpu: 1
    volumeMounts:
    - mountPath: /kpuworkshop
      name: fs-store
    - mountPath: /home/jovyan
      name: $volname
  restartPolicy: Never
  volumes:
  - name: fs-store
    flexVolume:
      driver: ceph.rook.io/rook
      fsType: ceph
      options:
        clusterNamespace: rook
        fsName: nautilusfs
        path: /kpuworkshop
        mountUser: kpuworkshop
        mountSecret: ceph-fs-secret
  - name: $volname
    persistentVolumeClaim:
      claimName: $volname
EOF
    sleep 5
else
    printf "Connecting to existing pod '$podname'.\n"
fi
while ! kubectl get pods --namespace=$namespace --field-selector=status.phase=Running 2> /dev/null | grep "$podname" -q ; do
    printf "Wait for pod...\n"
    sleep 5
done


########################
# Start Jupyter in pod #
########################
mlineno=$LINENO
printf "\n\n"
sed "$((mlineno-3)),$((mlineno-1))!d" $0 | colorize

jupyter_pid=42514251
(while ! ps -p "$jupyter_pid" > /dev/null ; do
    printf "Please wait (about 5 secs)\n"
    (kubectl exec --namespace=$namespace -it "$podname" bash <<EOF & echo $! >.tmp_${podname}_pid
jupyter notebook --ip='0.0.0.0'
EOF
    ) 2>&1 | tee .tmp_${podname}_logfile_jupyter > /dev/null &
    jupyter_pid=$(cat .tmp_${podname}_pid)
    sleep 5
    head -n -10 .tmp_${podname}_logfile_jupyter | tail -10
    if [ "$jupyter_pid" = "" ] ; then
	jupyter_pid=42514251
    fi
 done)
jupyter_pid=$(cat .tmp_${podname}_pid)
if ! [ -f .tmp_${podname}_logfile_jupyter ] ; then
    echo "Unexpected missing log file"
    exit 1
fi
jupyter_port=$(cat .tmp_${podname}_logfile_jupyter | sed -n "s|.*/($podname or 127.0.0.1):\([0-9]*\)/.*|\1|p" | head -1)
while [ "$jupyter_port" = "" ] ; do
    printf "Could not figure out the Jupyter port. Retrying...\n"
    sleep 5
    head -n -10 .tmp_${podname}_logfile_jupyter | tail -10
    jupyter_port=$(cat .tmp_${podname}_logfile_jupyter | sed -n "s|.*/($podname or 127.0.0.1):\([0-9]*\)/.*|\1|p" | head -1)
done
printf "\n"
printf "Jupyter PID=$jupyter_pid PORT=$jupyter_port\n"


########################################
# Start port forwarding (about 5 secs) #
########################################
mlineno=$LINENO
printf "\n\n"
sed "$((mlineno-3)),$((mlineno-1))!d" $0 | colorize

(kubectl port-forward "$podname" $local_port:$jupyter_port -n "$namespace" & echo $! >.tmp_${podname}_pid) 2>&1 | tee .tmp_${podname}_logfile_port_fw > /dev/null &
port_fw_pid=$(cat .tmp_${podname}_pid)
sleep 5
if ! ps -p "$port_fw_pid" > /dev/null ; then
    echo "Port forwarding ended unexpectedly"
    exit 1
fi
cat .tmp_${podname}_logfile_port_fw
printf "\n"
printf "Copy in your Web browser\n\t"
cat .tmp_${podname}_logfile_jupyter | sed -n "s|.*http://($podname or 127.0.0.1):[0-9]*/\(.*\)|http://localhost:$local_port/\1|p" | head -1
printf "\n"
printf "Port forwarding PID=$port_fw_pid FROM=$local_port TO=$jupyter_port\n"


###################################################
# Start bash session in pod (type 'exit' to quit) #
###################################################
mlineno=$LINENO
printf "\n\n"
sed "$((mlineno-3)),$((mlineno-1))!d" $0 | colorize

while ! kubectl exec --namespace=$namespace -it "$podname" bash ; do
    echo $?
    printf "Retry. Please wait (about 5 secs)\n"
    sleep 5
done


############################
# Stopping and cleaning up #
############################
mlineno=$LINENO
printf "\n\n"
sed "$((mlineno-3)),$((mlineno-1))!d" $0 | colorize

printf "Stop port forwarding\n"
kill -9 "$port_fw_pid"
printf "Stop Jupyter\n"
kill -9 "$jupyter_pid"
rm -f .tmp_${podname}_logfile_*
rm -f .tmp_${podname}_pid
#printf "Delete pod\n"
#kubectl delete pod --namespace=$namespace "$podname"