#!/bin/sh

set -euo pipefail

# Apply cgroup v2 hack
cgroup-v2-hack.sh

trap 'pkill -P $$' EXIT
FLYTE_TIMEOUT=${FLYTE_TIMEOUT:-600}

monitor() {
    while : ; do
        for pid in $@ ; do
            kill -0 $pid &> /dev/null || exit 1
        done

        sleep 1
    done
}

# Start docker daemon
echo "Starting Docker daemon..."
file="/var/run/docker.pid"
if [ -f "$file" ] ; then
    rm "$file"
fi
dockerd &> /var/log/dockerd.log &
DOCKERD_PID=$!
timeout "$FLYTE_TIMEOUT" sh -c "until docker info &> /dev/null; do sleep 1; done" || ( echo >&2 "Timed out while waiting for dockerd to start"; exit 1 )
echo "Done."

# Start k3s
echo "Starting k3s cluster..."
KUBERNETES_API_PORT=${KUBERNETES_API_PORT:-6443}
k3s server --docker --no-deploy=traefik --no-deploy=servicelb --no-deploy=local-storage --no-deploy=metrics-server --https-listen-port=${KUBERNETES_API_PORT} &> /var/log/k3s.log &
K3S_PID=$!
timeout "$FLYTE_TIMEOUT" sh -c "until k3s kubectl get node $HOSTNAME &> /dev/null; do sleep 1; done" || ( echo >&2 "Timed out while waiting for the Kubernetes cluster to start"; exit 1 )
k3s kubectl wait node $HOSTNAME --for condition=Ready --timeout ${FLYTE_TIMEOUT}s &> /dev/null || ( echo >&2 "Timed out while waiting for the Kubernetes cluster to be ready"; exit 1 )
echo "Done."

echo "Deploying Flyte..."
FLYTE_VERSION=${FLYTE_VERSION:-latest}
if [[ $FLYTE_VERSION = "latest" ]]
then
  FLYTE_VERSION=$(curl --silent "https://api.github.com/repos/flyteorg/flyte/releases/latest" | jq -r .tag_name)
fi

# Deploy flyte
helm repo add flyteorg https://flyteorg.github.io/flyte

echo "Deploying Flyte-deps..."
version=""
charts="/flyteorg/share/flyte-deps"

if [[ $FLYTE_TEST = "release" ]]
then
  version="--version $FLYTE_VERSION"
  charts="flyteorg/flyte-deps"
fi

if [[ $FLYTE_TEST = "local" ]]
then
  helm dep update $charts
fi

helm upgrade -n flyte --create-namespace flyte-deps $charts --kubeconfig /etc/rancher/k3s/k3s.yaml --install $version --set webhook.enabled=false,contour.enabled=true --wait

version=""
charts="/flyteorg/share/flyte-core"

if [[ $FLYTE_TEST = "release" ]]
then
  version="--version $FLYTE_VERSION"
  charts="flyteorg/flyte-core"
fi

if [[ $FLYTE_TEST = "local" ]]
then
  helm dep update $charts
fi

helm upgrade -n flyte --create-namespace flyte-core $charts --kubeconfig /etc/rancher/k3s/k3s.yaml --install $version -f /flyteorg/share/flyte-core/values-sandbox.yaml

wait-for-flyte.sh

# With flytectl sandbox --source flag, we mount the root volume to user source dir that will create helm & k8s cache specific directory.
# In Linux, These file belongs to root user that is different then current user
# In this case during fast serialization, Pyflyte will through error because of permission denied
rm -rf /root/.cache /root/.kube /root/.config

# Monitor running processes. Exit when the first process exits.
monitor ${DOCKERD_PID} ${K3S_PID}
