#!/usr/bin/env bash
# Daemonless image builder for RunPod serverless endpoints (crane-based; no
# Docker needed — runs on any RunPod pod that has the models on disk).
#
# Layer model (order matters — cheap-to-change stuff goes LAST):
#
#   plx1029/comfyui-qwen:v6                       (base: torch + ComfyUI + rgthree)
#     + custom-nodes layer + one layer PER model  -> :<endpoint>-models-<V>   (heavy, build rarely)
#       + serverless code layer + CMD mutate      -> :<endpoint>-<V>          (tiny, build often)
#
# Because the code layer sits on top, iterating on handler/workflows/params
# re-uploads only a few MB; the model layers are content-addressed and are
# never pushed twice.
#
# Usage (from the directory holding this script):
#   ./build.sh models qwen-edit-turbo v1            # build+push heavy models tag
#   ./build.sh code   qwen-edit-turbo v1 v1         # build+push final tag on models-v1
#   ./build.sh all    qwen-edit-turbo v1            # both
#
# Env: DOCKER_HUB_ACCESS_TOKEN (required), DOCKER_HUB_USER (default plx1029),
#      BASE_IMAGE (default plx1029/comfyui-qwen:v6),
#      DEST_REPO (default docker.io/<user>/comfyui-serverless),
#      COMFY_WORKSPACE (default /workspace).

set -euo pipefail

CMD_MODE="${1:?models|code|all}"; ENDPOINT="${2:?endpoint name}"; VER="${3:?version}"
MODELS_VER="${4:-$VER}"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
WS="${COMFY_WORKSPACE:-/workspace}"
USER="${DOCKER_HUB_USER:-plx1029}"
BASE="${BASE_IMAGE:-plx1029/comfyui-qwen:v6}"
REPO="${DEST_REPO:-docker.io/${USER}/comfyui-serverless}"
EP_DIR="$SCRIPT_DIR/endpoints/$ENDPOINT"
BUILD_DIR="${BUILD_DIR:-/root/build}"
THREADS=$(( $(nproc) > 64 ? 64 : $(nproc) ))

[ -f "$EP_DIR/endpoint.json" ] || { echo "no endpoint.json in $EP_DIR"; exit 1; }
command -v crane >/dev/null || { echo "crane not installed"; exit 1; }
command -v pigz  >/dev/null || { echo "pigz not installed";  exit 1; }

echo "$DOCKER_HUB_ACCESS_TOKEN" | crane auth login index.docker.io -u "$USER" --password-stdin

mkdir -p "$BUILD_DIR"
# .git excluded: images are immutable, git history is dead weight (RES4LYF's
# alone is 260 MB) and Manager-updates don't apply to serverless workers
TAR_OPTS=(--owner=0 --group=0 --exclude='__pycache__' --exclude='*.pyc' --exclude='.git')

layer_from_paths() {  # $1 = output tgz, rest = paths relative to /
    local out="$1"; shift
    [ -s "$out" ] && { echo "  (cached) $out"; return; }
    tar "${TAR_OPTS[@]}" -C / -cf - "$@" | pigz -p "$THREADS" -1 > "$out.tmp"
    mv "$out.tmp" "$out"
    echo "  $(du -h "$out" | cut -f1) $out"
}

build_models() {
    local tag="$REPO:$ENDPOINT-models-$VER"
    echo "== building $tag on $BASE"
    local layers=()

    # custom nodes the endpoint needs that aren't already in the base image
    # (BASE_HAS_NODES lists packs baked into $BASE; rgthree ships in v6)
    local base_has=" ${BASE_HAS_NODES:-rgthree-comfy} "
    local nodes n paths=()
    nodes=$(python3 -c "import json;print(' '.join(json.load(open('$EP_DIR/endpoint.json'))['custom_nodes']))")
    for n in $nodes; do
        case "$base_has" in *" $n "*) continue ;; esac
        [ -d "$WS/ComfyUI/custom_nodes/$n" ] || { echo "missing custom node $n"; exit 1; }
        paths+=("workspace/ComfyUI/custom_nodes/$n")
    done
    if [ "${#paths[@]}" -gt 0 ]; then
        echo "-- custom nodes layer: ${paths[*]}"
        layer_from_paths "$BUILD_DIR/$ENDPOINT-nodes.tgz" "${paths[@]}"
        layers+=(-f "$BUILD_DIR/$ENDPOINT-nodes.tgz")
    fi

    # one layer per model file/dir (content-addressed => pushed exactly once, ever)
    while read -r m; do
        [ -e "$WS/ComfyUI/$m" ] || { echo "missing model $WS/ComfyUI/$m — download it first"; exit 1; }
        local tgz="$BUILD_DIR/model-$(basename "$m").tgz"
        echo "-- model layer: $m"
        layer_from_paths "$tgz" "workspace/ComfyUI/$m"
        layers+=(-f "$tgz")
    done < <(python3 -c "import json;print('\n'.join(json.load(open('$EP_DIR/endpoint.json'))['models']))")

    # extra prebuilt layers (e.g. a pip-deps delta captured on the build pod)
    while read -r x; do
        [ -z "$x" ] && continue
        [ -s "$BUILD_DIR/$x" ] || { echo "missing extra layer $BUILD_DIR/$x"; exit 1; }
        echo "-- extra layer: $x"
        layers+=(-f "$BUILD_DIR/$x")
    done < <(python3 -c "import json;print('\n'.join(json.load(open('$EP_DIR/endpoint.json')).get('extra_layers',[])))")

    crane append -b "$BASE" "${layers[@]}" -t "$tag"
    echo "== pushed $tag"
}

build_code() {
    local models_tag="$REPO:$ENDPOINT-models-$MODELS_VER"
    local tag="$REPO:$ENDPOINT-$VER"
    echo "== building $tag on $models_tag"

    local stage="$BUILD_DIR/stage-code"
    rm -rf "$stage"
    mkdir -p "$stage/opt/serverless/src" "$stage/opt/serverless/lib" \
             "$stage/workspace/ComfyUI/comfy"
    cp -r "$SCRIPT_DIR"/handler.py "$SCRIPT_DIR"/start.sh "$SCRIPT_DIR"/endpoints \
          "$SCRIPT_DIR"/tests "$stage/opt/serverless/src/"
    # only THIS endpoint's config goes into the image
    find "$stage/opt/serverless/src/endpoints" -mindepth 1 -maxdepth 1 \
         ! -name "$ENDPOINT" -exec rm -rf {} +
    # vendored python deps for the handler (runpod sdk); build them if absent
    if [ ! -d /opt/serverless/lib/runpod ]; then
        pip install -q --target /opt/serverless/lib runpod
    fi
    cp -r /opt/serverless/lib/. "$stage/opt/serverless/lib/"
    # local core patches (ComfyUI #14573 aimdo fix) ride along with the code
    cp "$WS/ComfyUI/comfy/model_management.py" "$stage/workspace/ComfyUI/comfy/"
    chmod +x "$stage/opt/serverless/src/start.sh"

    tar "${TAR_OPTS[@]}" --exclude='local-out-*' -C "$stage" -cf - . \
        | pigz -p "$THREADS" -1 > "$BUILD_DIR/code.tgz"
    crane append -b "$models_tag" -f "$BUILD_DIR/code.tgz" -t "$tag"
    crane mutate "$tag" --cmd "/opt/serverless/src/start.sh" -t "$tag"
    echo "== pushed $tag"
}

case "$CMD_MODE" in
    models) build_models ;;
    code)   build_code ;;
    all)    build_models; MODELS_VER="$VER" build_code ;;
    *) echo "unknown mode $CMD_MODE"; exit 1 ;;
esac