#!/bin/bash
# Step 0 of the image/video plan: try ComfyUI by hand on this machine and measure it, before any code is written.
#
# Download it once, then run the commands one at a time:
#   curl -fsSL @@BASE@@/lab/comfyui-step0.sh -o ~/comfy-step0.sh
#   bash ~/comfy-step0.sh check
#
#   check          look only: CPU, memory, disk, Docker, what is loaded right now. Changes nothing.
#   prepare        pull the NVIDIA PyTorch image (about 20 GB) and fetch ComfyUI (its latest release tag)
#   models flux    FLUX.1-schnell fp8 (17.2 GB)
#   models wan     Wan2.2 TI2V-5B + VAE + text encoder (18 GB)
#   start          start ComfyUI in a container; it listens on 127.0.0.1 only (see the SSH tunnel it prints)
#   watch          print memory / GPU / CPU every 5 s into measure.log (run it in a second terminal while generating)
#   logs           follow ComfyUI's log ("Prompt executed in N seconds" is the time you want)
#   report         one block to paste back: versions, times, peak memory, what the Owner Guard did
#   stop           stop the container
#   clean          remove the container, ComfyUI and its Python packages (the models stay);  clean --all  removes everything
#
# Everything lives in $WXCH_LAB_DIR (default ~/comfy-lab). Nothing is installed on the system, no port is opened to the
# network, the container is not privileged and has no access to Docker or to your other files. Needs Docker (sudo is used
# for it if your user is not in the docker group) and the NVIDIA container runtime, which vLLM already needs.
set -uo pipefail
LAB="${WXCH_LAB_DIR:-$HOME/comfy-lab}"
IMAGE="${WXCH_LAB_IMAGE:-nvcr.io/nvidia/pytorch:26.09-py3}"
NAME=wxch-comfy-lab
PORT=8188

say()  { printf '\033[1m[lab]\033[0m %s\n' "$*"; }
warn() { printf '\033[33m[lab] warning:\033[0m %s\n' "$*" >&2; }
die()  { printf '\033[31m[lab] error:\033[0m %s\n' "$*" >&2; exit 1; }
have() { command -v "$1" >/dev/null 2>&1; }
D() { if docker info >/dev/null 2>&1; then docker "$@"; else sudo docker "$@"; fi; }
gbfree() { df --output=avail -BG "$1" 2>/dev/null | tail -1 | tr -dc '0-9'; }

need_docker() { have docker || die "Docker is not installed on this machine"; D info >/dev/null 2>&1 || die "cannot talk to Docker (is it running?)"; }

cmd_check() {
  echo "machine:   $(uname -m), $(. /etc/os-release 2>/dev/null; echo "${PRETTY_NAME:-linux}"), kernel $(uname -r)"
  echo "memory:    $(free -g | awk '/^Mem:/ {print "total " $2 " GB, used " $3 " GB, available " $7 " GB"}')"
  mkdir -p "$(dirname "$LAB")" 2>/dev/null
  echo "disk:      $(gbfree "$(dirname "$LAB")") GB free where the lab lives ($LAB); about 80 GB are needed for everything"
  if have nvidia-smi; then echo "gpu:       $(nvidia-smi --query-gpu=name,driver_version --format=csv,noheader 2>/dev/null | head -1)"; else echo "gpu:       nvidia-smi not found"; fi
  if have docker; then
    if D info >/dev/null 2>&1; then
      echo "docker:    $(D info --format '{{.ServerVersion}}, runtimes: {{range $k, $v := .Runtimes}}{{$k}} {{end}}' 2>/dev/null)"
      D image inspect "$IMAGE" >/dev/null 2>&1 && echo "image:     $IMAGE is already pulled" || echo "image:     $IMAGE not pulled yet (prepare does it)"
      D ps --filter "name=$NAME" --format 'container: {{.Names}} {{.Status}}' 2>/dev/null | grep . || echo "container: none running"
    else
      echo "docker:    installed, but this user cannot use it without sudo"
    fi
  else
    echo "docker:    not installed"
  fi
  echo "ollama:    $(curl -fsS -m 3 http://127.0.0.1:11434/api/ps 2>/dev/null | python3 -c 'import json,sys; m=[x["name"] for x in json.load(sys.stdin).get("models",[])]; print(", ".join(m) if m else "no model loaded")' 2>/dev/null || echo 'not reachable')"
  ss -ltn 2>/dev/null | grep -q ":$PORT " && echo "port:      $PORT is already in use" || echo "port:      $PORT is free"
  [ -d "$LAB" ] && echo "lab:       $LAB exists ($(du -sh "$LAB" 2>/dev/null | cut -f1))" || echo "lab:       $LAB does not exist yet"
}

cmd_prepare() {
  need_docker
  have git || die "git is needed"
  have curl || die "curl is needed"
  mkdir -p "$LAB"; : > "$LAB/.wxch-lab"
  [ "$(gbfree "$LAB")" -ge 30 ] || die "less than 30 GB free in $LAB"
  say "pulling $IMAGE (about 20 GB; if the connection drops it resumes when you run this again)"
  n=0
  until D pull "$IMAGE"; do
    n=$((n+1)); [ "$n" -lt 5 ] || die "could not pull the image after $n attempts: check the connection and run prepare again"
    say "pull interrupted (attempt $n of 5), retrying in $((n*15)) s"; sleep $((n*15))
  done
  if [ ! -d "$LAB/ComfyUI/.git" ]; then
    tag="$(git ls-remote --tags --sort=-v:refname https://github.com/comfyanonymous/ComfyUI 'v*' 2>/dev/null | head -1 | sed 's#.*refs/tags/##')"
    [ -n "$tag" ] || die "cannot read ComfyUI's release tags from GitHub"
    say "fetching ComfyUI $tag"
    git clone -q --depth 1 --branch "$tag" https://github.com/comfyanonymous/ComfyUI "$LAB/ComfyUI" || die "git clone failed"
  fi
  # the container brings its own PyTorch built for this GPU: never let pip replace it
  grep -viE '^(torch|torchvision|torchaudio)([=<>~! ]|$)' "$LAB/ComfyUI/requirements.txt" > "$LAB/ComfyUI/requirements.lab.txt"
  mkdir -p "$LAB/pyuser" "$LAB/ComfyUI/models/checkpoints" "$LAB/ComfyUI/models/diffusion_models" "$LAB/ComfyUI/models/vae" "$LAB/ComfyUI/models/text_encoders"
  say "ready: ComfyUI $(git -C "$LAB/ComfyUI" describe --tags 2>/dev/null) in $LAB. Next: models flux (and/or models wan), then start"
}

# the size the server says the file has (the last Content-Length after the redirects)
remote_size() { curl -fsSIL -m 30 "$1" 2>/dev/null | tr -d '\r' | awk 'tolower($1)=="content-length:" {v=$2} END {print v+0}'; }

dl() { # url dest
  url="$1"; dest="$2"
  size="$(remote_size "$url")"
  [ "${size:-0}" -gt 1000000 ] || die "cannot read the size of $(basename "$dest") from the server (is the network up?)"
  if [ -f "$dest" ] && [ "$(stat -c %s "$dest")" = "$size" ]; then say "already have $(basename "$dest")"; return 0; fi
  say "downloading $(basename "$dest") ($(( size / 1000000 )) MB); it resumes if interrupted"
  n=0
  until curl -fL --retry 3 --retry-delay 5 -C - -o "$dest.part" "$url"; do
    n=$((n+1)); [ "$n" -lt 6 ] || die "download failed after $n attempts: run the same command again to resume"
    sleep 10
  done
  [ "$(stat -c %s "$dest.part")" = "$size" ] || die "$(basename "$dest") has the wrong size ($(stat -c %s "$dest.part") of $size bytes): delete $dest.part and try again"
  mv "$dest.part" "$dest"
}

cmd_models() {
  [ -d "$LAB/ComfyUI" ] || die "run prepare first"
  m="$LAB/ComfyUI/models"
  case "${1:-}" in
    flux)
      [ "$(gbfree "$LAB")" -ge 20 ] || die "less than 20 GB free in $LAB"
      dl "https://huggingface.co/Comfy-Org/flux1-schnell/resolve/main/flux1-schnell-fp8.safetensors" "$m/checkpoints/flux1-schnell-fp8.safetensors" ;;
    wan)
      [ "$(gbfree "$LAB")" -ge 20 ] || die "less than 20 GB free in $LAB"
      base="https://huggingface.co/Comfy-Org/Wan_2.2_ComfyUI_Repackaged/resolve/main/split_files"
      dl "$base/diffusion_models/wan2.2_ti2v_5B_fp16.safetensors" "$m/diffusion_models/wan2.2_ti2v_5B_fp16.safetensors"
      dl "$base/vae/wan2.2_vae.safetensors" "$m/vae/wan2.2_vae.safetensors"
      dl "$base/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors" "$m/text_encoders/umt5_xxl_fp8_e4m3fn_scaled.safetensors" ;;
    *) die "usage: models flux | models wan" ;;
  esac
  say "done"
}

cmd_start() {
  need_docker
  [ -f "$LAB/ComfyUI/requirements.lab.txt" ] || die "run prepare first"
  D rm -f "$NAME" >/dev/null 2>&1
  say "starting ComfyUI (the first start installs its Python packages and takes a minute or two)"
  # runs as your user, not root; models and output stay in $LAB; only 127.0.0.1 can reach the port
  D run -d --name "$NAME" --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 \
    --user "$(id -u):$(id -g)" -e HOME=/work -e PYTHONUSERBASE=/work/pyuser -p "127.0.0.1:$PORT:8188" \
    -v "$LAB:/work" -w /work/ComfyUI "$IMAGE" \
    bash -c 'pip install --user -q -r requirements.lab.txt && exec python main.py --listen 0.0.0.0 --port 8188 --disable-auto-launch' >/dev/null || die "could not start the container"
  i=0
  until curl -fsS -m 3 "http://127.0.0.1:$PORT/system_stats" >/tmp/wxch-comfy-stats.json 2>/dev/null; do
    i=$((i+1)); [ "$i" -lt 90 ] || { warn "ComfyUI did not answer in 5 minutes: look at '$0 logs'"; return 1; }
    D ps --filter "name=$NAME" --filter status=running -q | grep -q . || { warn "the container stopped: look at '$0 logs'"; return 1; }
    sleep 3
  done
  say "ComfyUI is up. What it sees:"
  python3 -c 'import json; d=json.load(open("/tmp/wxch-comfy-stats.json")); s=d.get("system",{}); print("  ComfyUI", s.get("comfyui_version","?"), "· torch", s.get("pytorch_version","?")); [print("  device:", x.get("name"), "· vram_total", x.get("vram_total"), "· torch_vram_total", x.get("torch_vram_total")) for x in d.get("devices",[])]' 2>/dev/null
  echo
  echo "To open it from your own computer, run this THERE (not here), then browse http://127.0.0.1:$PORT :"
  echo "  ssh -N -L $PORT:127.0.0.1:$PORT $(id -un)@$(hostname -I 2>/dev/null | awk '{print $1}')"
  echo
  echo "In the page: Workflow > Browse Templates > search 'Flux' (Flux Schnell) or 'Wan2.2' (5B video). Run each twice: the first run includes loading the model."
}

cmd_watch() {
  f="$LAB/measure.log"
  [ -f "$f" ] || echo "time,mem_used_mb,mem_avail_mb,gpu_util_pct,gpu_temp_c,gpu_power_w,container_mem,container_cpu" > "$f"
  say "writing to $f every 5 s (Ctrl-C to stop)"
  while true; do
    t="$(date +%T)"
    read -r used avail < <(free -m | awk '/^Mem:/ {print $3, $7}')
    g="$(nvidia-smi --query-gpu=utilization.gpu,temperature.gpu,power.draw --format=csv,noheader,nounits 2>/dev/null | head -1 | tr -d ' ')"
    c="$(D stats --no-stream --format '{{.MemUsage}};{{.CPUPerc}}' "$NAME" 2>/dev/null | sed 's#;#,#; s# /[^,]*##')"
    line="$t,$used,$avail,${g:-n/a,n/a,n/a},${c:-none,none}"
    echo "$line" | tee -a "$f"
    sleep 5
  done
}

cmd_logs() { need_docker; D logs -f --tail 80 "$NAME"; }

cmd_report() {
  echo "===== comfyui step 0 report ($(date '+%F %T')) ====="
  cmd_check
  echo "--- versions"
  D exec "$NAME" python -c 'import torch; print("torch", torch.__version__, "cuda", torch.version.cuda, "| device", torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))' 2>&1 | tail -2
  echo "--- times and problems (from ComfyUI's log)"
  D logs "$NAME" 2>&1 | grep -E "Prompt executed in|Requested to load|loaded completely|out of memory|OutOfMemory|CUDA error|no kernel image|Traceback|Error" | tail -30
  echo "--- memory over time (peak of used memory, MB)"
  if [ -f "$LAB/measure.log" ]; then
    awk -F, 'NR>1 && $2+0>m {m=$2+0; t=$1} END {print "peak used:", m, "MB at", t}' "$LAB/measure.log"
    awk -F, 'NR>1 {if (min==""||$3+0<min) min=$3+0} END {print "lowest available:", min, "MB"}' "$LAB/measure.log"
  else
    echo "no measure.log (run 'watch' in a second terminal while generating)"
  fi
  echo "--- what the Owner Guard did"
  journalctl -u wxch-agent --since "-3 hours" --no-pager 2>/dev/null | grep "owner guard" | tail -8 || echo "(cannot read the agent log as this user)"
  echo "===== end ====="
}

cmd_stop() { need_docker; D rm -f "$NAME" >/dev/null 2>&1 && say "stopped"; }

cmd_clean() {
  [ -f "$LAB/.wxch-lab" ] || die "$LAB was not made by this script: refusing to delete anything"
  need_docker
  D rm -f "$NAME" >/dev/null 2>&1
  if [ "${1:-}" = "--all" ]; then
    rm -rf "$LAB" && say "removed $LAB (including the models)"
    D rmi "$IMAGE" >/dev/null 2>&1 && say "removed the image"
  else
    rm -rf "$LAB/ComfyUI/.git" "$LAB/pyuser" "$LAB/measure.log"
    find "$LAB/ComfyUI" -mindepth 1 -maxdepth 1 ! -name models -exec rm -rf {} + 2>/dev/null
    say "removed the container, ComfyUI and its packages; the models in $LAB/ComfyUI/models stay ('clean --all' removes them too)"
  fi
}

case "${1:-}" in
  check) cmd_check ;;
  prepare) cmd_prepare ;;
  models) shift; cmd_models "$@" ;;
  start) cmd_start ;;
  watch) cmd_watch ;;
  logs) cmd_logs ;;
  report) cmd_report ;;
  stop) cmd_stop ;;
  clean) shift; cmd_clean "$@" ;;
  *) sed -n '2,/^set -uo/p' "$0" | sed '$d' | sed 's/^# \{0,1\}//'; [ -z "${1:-}" ] || exit 2 ;;
esac