#!/usr/bin/env bash
# Source inside a Slurm batch allocation, then call start_study_docker ROOT.
set -euo pipefail

start_study_docker() {
    local root="${1:?study root required}"
    : "${SLURM_JOB_ID:?Rootless Docker must run inside a Slurm allocation}"
    export PATH="$root/runtime/bin:$root/runtime/node-v22.21.0-linux-x64/bin:$PATH"
    study_docker_tmp="$(mktemp -d "/tmp/beam-pb-${SLURM_JOB_ID}.XXXXXXXX")"
    export XDG_RUNTIME_DIR="$study_docker_tmp/run"
    export XDG_CONFIG_HOME="$study_docker_tmp/config"
    export DOCKER_CONFIG="$study_docker_tmp/docker-config"
    export DOCKER_HOST="unix://$XDG_RUNTIME_DIR/docker.sock"
    if [[ -f "$root/runtime/evaluator-wheelhouse/manifest.json" ]]; then
        export PROGRAMBENCH_EVALUATOR_WHEELHOUSE="$root/runtime/evaluator-wheelhouse"
    fi
    mkdir -p "$XDG_RUNTIME_DIR" "$XDG_CONFIG_HOME" "$DOCKER_CONFIG" "$study_docker_tmp/data" "$root/logs"
    chmod 700 "$study_docker_tmp" "$XDG_RUNTIME_DIR" "$DOCKER_CONFIG"
    export DOCKERD_ROOTLESS_ROOTLESSKIT_NET=slirp4netns
    export DOCKERD_ROOTLESS_ROOTLESSKIT_PORT_DRIVER=slirp4netns
    dockerd-rootless.sh --host="$DOCKER_HOST" --data-root="$study_docker_tmp/data" \
      --exec-root="$study_docker_tmp/exec" --pidfile="$study_docker_tmp/dockerd.pid" \
      --exec-opt native.cgroupdriver=cgroupfs --storage-driver=fuse-overlayfs \
      --iptables=false --bridge=none > "$root/logs/rootless-daemon-$SLURM_JOB_ID.log" 2>&1 &
    study_docker_pid=$!
    for ((attempt=0; attempt<90; attempt++)); do
        if docker info >/dev/null 2>&1; then
            docker info --format '{{json .}}' > "$root/logs/docker-info-$SLURM_JOB_ID.json"
            printf 'Docker ready: job=%s node=%s socket=%s\n' "$SLURM_JOB_ID" "$(hostname)" "$DOCKER_HOST"
            return 0
        fi
        if ! kill -0 "$study_docker_pid" 2>/dev/null; then
            tail -60 "$root/logs/rootless-daemon-$SLURM_JOB_ID.log"
            return 1
        fi
        sleep 1
    done
    tail -60 "$root/logs/rootless-daemon-$SLURM_JOB_ID.log"
    return 1
}

stop_study_docker() {
    if [[ -n "${study_docker_pid:-}" ]]; then
        kill "$study_docker_pid" 2>/dev/null || true
        wait "$study_docker_pid" 2>/dev/null || true
    fi
    # The job-scoped /tmp data holds shifted UID files. Slurm/node cleanup owns it;
    # do not invoke a privileged cleanup or remove unrelated runtime directories.
}

ensure_study_images() {
    local root="${1:?study root required}"
    local python="$root/runtime/programbench-venv/bin/python"
    local cache="$root/runtime/programbench-images.tar.gz"
    local -a image_refs image_aliases
    mapfile -t image_refs < <("$python" -c 'import json,sys; print("\n".join(t["task_cleanroom_image"]["immutable_reference"] for t in json.load(open(sys.argv[1]))["tasks"]))' "$root/docs/programbench-five-task-manifest.json")
    if [[ -s "$cache" ]]; then
        printf 'Loading cached ProgramBench image layers\n'
        gzip -dc "$cache" | docker image load > "$root/logs/image-load-$SLURM_JOB_ID.log"
    fi
    image_aliases=()
    local index=0
    for image in "${image_refs[@]}"; do
        printf 'Ensuring pinned ProgramBench image %s/%s\n' "$((index + 1))" "${#image_refs[@]}"
        # Docker save/load does not preserve repository digest associations. This pull
        # restores them; loaded layers are reused and the requested digest is immutable.
        docker pull "$image" >> "$root/logs/image-pull-$SLURM_JOB_ID.log" 2>&1
        local alias="study-pb-cache/task$index:${image##*:}"
        docker tag "$image" "$alias"
        image_aliases+=("$alias")
        index=$((index + 1))
    done
    if [[ ! -s "$cache" ]]; then
        printf 'Caching image layers for the study allocation\n'
        docker image save "${image_aliases[@]}" | gzip -1 > "$cache.tmp-$SLURM_JOB_ID"
        mv "$cache.tmp-$SLURM_JOB_ID" "$cache"
        sha256sum "$cache" > "$cache.sha256"
    fi
}