#!/usr/bin/env bash
# KidCAD update (V1.2.3) for Linux and Mac. Finds your KidCAD folder and the newest kidcad*.zip by itself.
# Your family's things (the data folder), the AI/figure environments and your own printer profiles are kept.
# Do NOT run this while a USB print is running: KidCAD restarts at the end and that stops a USB print.
set -uo pipefail
say()  { printf '\n\033[1;35m==> %s\033[0m\n' "$*"; }
warn() { printf '\033[1;33m[!] %s\033[0m\n' "$*"; }
die()  { printf '\033[1;31m[x] %s\033[0m\n' "$*"; exit 1; }

# ---- 1. the newest KidCAD zip for this computer (Downloads, home, Desktop)
if [[ "$(uname -s)" == "Darwin" ]]; then OTHER='(windows|linux)'; INSTALLER=install_mac.sh; else OTHER='(windows|mac)'; INSTALLER=install.sh; fi
ZIP=""
while IFS= read -r z; do  # newest first
  if [[ -z "$ZIP" || "$z" -nt "$ZIP" ]]; then ZIP="$z"; fi
done < <(find "$HOME/Downloads" "$HOME" "$HOME/Desktop" -maxdepth 2 -iname 'kidcad*.zip' -type f 2>/dev/null | grep -viE "_${OTHER}\.zip$" | grep -viE '_source\.zip$')
[[ -n "$ZIP" ]] || die "No kidcad*.zip found in Downloads, Desktop or your home folder."
say "Found this download: $ZIP  ($(date -r "$ZIP" '+%b %d %H:%M'), $(du -h "$ZIP" | cut -f1))"
read -r -p "Install this one? Only say yes if you downloaded it yourself from kidcad.xyz or GitHub. (y/N): " okzip
[[ "$okzip" =~ ^[Yy] ]] || die "Nothing changed."

# ---- 2. the installed KidCAD folder (the one with app/server.py and its .venv)
DIRS=()
while IFS= read -r d; do DIRS+=("$d"); done < <(find "$HOME" -maxdepth 5 \( -name Library -o -name .Trash -o -name .cache \
    -o -name .venv -o -name 'kidcad_update_tmp*' -o -name node_modules \) -prune -o -path '*/app/server.py' -print 2>/dev/null \
    | sed 's#/app/server.py$##' \
    | while read -r d; do [[ -x "$d/.venv/bin/python" ]] && echo "$d"; done)
if [[ ${#DIRS[@]} -eq 0 ]]; then die "Couldn't find an installed KidCAD folder (one with app/server.py and .venv)."; fi
if [[ ${#DIRS[@]} -gt 1 ]]; then
  echo "More than one KidCAD folder found:"
  for i in "${!DIRS[@]}"; do echo "  $((i + 1))) ${DIRS[$i]}"; done
  read -r -p "Which one should be updated? (number): " n
  [[ "$n" =~ ^[0-9]+$ && $n -ge 1 && $n -le ${#DIRS[@]} ]] || die "No folder picked. Nothing changed."
  K="${DIRS[$((n - 1))]}"
else
  K="${DIRS[0]}"
fi
say "Updating KidCAD in: $K"

# ---- 3. never during a USB print
read -r -p "Is the printer printing over USB right now? (y/N): " busy
if [[ "$busy" =~ ^[Yy] ]]; then die "Run this again after the print is done. Nothing changed."; fi

# ---- 4. unpack and check the download
TMP="$HOME/kidcad_update_tmp_$$"
rm -rf "$TMP"; mkdir -p "$TMP"
unzip -q "$ZIP" -d "$TMP" || die "The zip could not be opened. Download it again."
SRC=$(dirname "$(find "$TMP" -maxdepth 3 -path '*/app/server.py' | head -1)")
SRC=$(dirname "$SRC")
[[ -f "$SRC/app/server.py" && -f "$SRC/run.sh" ]] || { rm -rf "$TMP"; die "That zip doesn't look like KidCAD."; }
echo "New version: $(cat "$SRC/VERSION" 2>/dev/null || echo '?')   Installed: $(cat "$K/VERSION" 2>/dev/null || echo 'before V1')"

# ---- 5. safety copy of the program and the family's data (only the newest copy is kept; the AI programs are
#         not copied, they can be installed again)
BK="$HOME/kidcad_backups"; STAMP=$(date +%Y%m%d_%H%M%S)
mkdir -p "$BK"
FREE_GB=$(df -Pk "$HOME" | awk 'NR==2 {print int($4 / 1048576)}')
if [[ "${FREE_GB:-0}" -lt 10 ]]; then
  warn "Less than 10 GB free on the disk (${FREE_GB} GB): no safety copy this time. Your last safety copy in $BK is kept."
  warn "To free space: KidCAD -> Grown-ups -> Storage -> Clean up now."
else
  say "Safety copy -> $BK"
  ok=1
  tar -czf "$BK/program_$STAMP.tar.gz" -C "$K" --exclude=./data --exclude=./.venv --exclude=./engines/.venv-figures \
      --exclude=./engines/Hunyuan3D-2 --exclude=./engines/.venv-triposg --exclude=./engines/TripoSG \
      --exclude='*/__pycache__' . && echo "program saved" || ok=0
  if [[ -d "$K/data" ]]; then
    tar -czf "$BK/data_$STAMP.tar.gz" -C "$K" --exclude=./data/tmp --exclude='*.gcode' data && echo "family data saved" || ok=0
  fi
  if [[ $ok -eq 1 ]]; then  # keep only this newest safety copy
    find "$BK" -maxdepth 1 -type f \( -name 'program_*.tar.gz' -o -name 'data_*.tar.gz' \) \
         ! -name "program_$STAMP.tar.gz" ! -name "data_$STAMP.tar.gz" -delete
    echo "older safety copies removed (the newest one is kept)"
  else
    warn "The safety copy had a problem, so the older safety copies were kept."
  fi
fi

# ---- 6. put the new program in place (data, environments and your printer edits are kept)
say "Copying the new KidCAD files"
if command -v rsync >/dev/null; then
  rsync -a --exclude '/data' --exclude '/.venv' --exclude '/engines/.venv-figures' --exclude '/engines/Hunyuan3D-2' \
        --exclude '/app/printers/' --exclude '__pycache__' "$SRC/" "$K/"
else
  (cd "$SRC" && find . -type f -not -path './data/*' -not -path './app/printers/*' -not -path '*/__pycache__/*' \
     -exec install -D -m 644 {} "$K/{}" \;)
fi
mkdir -p "$K/app/printers"
cp -n "$SRC"/app/printers/*.json "$K/app/printers/" 2>/dev/null   # new printers added, your old ones untouched
chmod +x "$K"/*.sh
rm -rf "$TMP"

# ---- 7. Python packages
say "Python packages"
"$K/.venv/bin/python" -m pip install -q --upgrade -r "$K/requirements.txt" || warn "pip had a problem (see above)."
if [[ -x "$K/engines/.venv-figures/bin/python" ]]; then
  say "Figure maker packages (keeps the versions that work)"
  "$K/engines/.venv-figures/bin/python" -m pip install -q "transformers<5" "huggingface-hub<1.0" diffusers \
    || warn "Figure maker packages had a problem; drawings and AI builds still work."
fi

# ---- 7b. optional: the TripoSG figure engine (MIT license, about 8 GB of graphics memory)
if command -v nvidia-smi >/dev/null && [[ ! -d "$K/engines/.venv-triposg" ]]; then
  read -r -p "Add the TripoSG figure engine now? It downloads several GB and takes a while. (y/N): " tsg
  if [[ "$tsg" =~ ^[Yy] ]]; then (cd "$K" && "./$INSTALLER" --triposg) || warn "TripoSG didn't install; the other figure models still work."; fi
fi

# ---- 8. check-up and restart
say "Check-up"
(cd "$K" && .venv/bin/python -m app.doctor) || true
say "Restarting KidCAD"
mkdir -p "$K/data/logs"
cd "$K"
if command -v setsid >/dev/null; then setsid nohup ./run.sh --restart >"$K/data/logs/restart.log" 2>&1 < /dev/null &
else nohup ./run.sh --restart >"$K/data/logs/restart.log" 2>&1 < /dev/null & fi
disown 2>/dev/null || true
sleep 6
if curl -s --max-time 5 "http://127.0.0.1:${KIDCAD_PORT:-7860}/api/auth" >/dev/null; then
  say "Done! KidCAD $(cat "$K/VERSION") is running: http://localhost:${KIDCAD_PORT:-7860}"
else
  warn "KidCAD didn't answer yet. Give it a few seconds, or start it from the desktop icon."
  echo "If it still won't start, the old program is in $BK/program_$STAMP.tar.gz"
fi