#!/usr/bin/env bash
# Copy code and the bundle to the pod, or bring runs back (run on this host; classifier-plan §2.7, §3.2, §3.3).
#   sync.sh up HOST PORT BUNDLE_DIR      training/ -> /workspace/code/training, bundle -> /root/data/bundle
#   sync.sh down HOST PORT BUNDLE_NAME [--final]
#                                        /workspace/runs -> ~/naver_psy_models/<bundle>/runs/ (after every stage);
#                                        --final also requires the pod's CLEAN marker (run.sh cleanup) before
#                                        the pod may be terminated
# SSH over the pod's exposed TCP port (the ssh.runpod.io proxy cannot rsync) with only the RunPod key
# (SSH_KEY, default ~/.ssh/naver_psy_runpod) and IdentitiesOnly. Uploads use -rl --checksum (the volume refuses
# chown and setting times). The bundle goes to the container disk only, never the network volume.
set -euo pipefail
HERE=$(cd "$(dirname "$0")/.." && pwd)              # .../training
REPO=$(cd "$HERE/.." && pwd)
dir=${1:-}; host=${2:-}; port=${3:-}
[ -n "$dir" ] && [ -n "$host" ] && [ -n "$port" ] || { sed -n 2,7p "$0"; exit 2; }
KEY=${SSH_KEY:-$HOME/.ssh/naver_psy_runpod}
SSH="ssh -p $port -i $KEY -o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new"
case "$dir" in
  up)
    bundle=${4:?bundle directory}
    [ -f "$bundle/manifest.json" ] || { echo "not a bundle: $bundle"; exit 1; }
    tree=$(cd "$HERE" && find . -type f ! -path '*/__pycache__/*' ! -name '.source_commit' ! -name '*.pyc' -print0 \
      | sort -z | xargs -0 sha256sum | sha256sum | cut -c1-16)
    commit=$(git -C "$REPO" rev-parse HEAD)
    [ -z "$(git -C "$REPO" status --porcelain -- "$HERE")" ] || commit="$commit-dirty"
    echo "$commit tree:$tree" > "$HERE/.source_commit"
    $SSH "root@$host" "mkdir -p /workspace/code /workspace/runs /root/data && chmod 700 /root/data"
    rsync -rlz --checksum --delete --exclude __pycache__ -e "$SSH" "$HERE/" "root@$host:/workspace/code/training/"
    rsync -rlz --checksum --chmod=D700,F600 -e "$SSH" "$bundle/" "root@$host:/root/data/bundle/"
    # Robustness inputs built locally (training.robustness build) hold post text too: container disk only.
    if [ -n "${ROBUSTNESS:-}" ]; then
      rsync -rlz --checksum --chmod=D700,F600 -e "$SSH" "$ROBUSTNESS/" "root@$host:/root/data/bundle/robustness/"
    fi
    echo "shipped $(cat "$HERE/.source_commit") and $(basename "$bundle")" ;;
  down)
    name=${4:?bundle name}
    LOCAL=${LOCAL_RUNS:-$HOME/naver_psy_models/$name}
    mkdir -p "$LOCAL"
    rsync -rltz -e "$SSH" --exclude 'last/trainer.pt' "root@$host:/workspace/runs/" "$LOCAL/runs/"
    for f in check_env.json selftest.json nvidia-smi.txt pip-freeze.txt; do
      rsync -rltz -e "$SSH" "root@$host:/workspace/$f" "$LOCAL/" 2>/dev/null || true
    done
    if [ "${5:-}" = --final ]; then
      $SSH "root@$host" "test -f /workspace/CLEAN && test ! -e /root/data" \
        || { echo "NOT CLEAN: run 'run.sh cleanup' on the pod before terminating"; exit 1; }
      echo "CLEAN verified; the pod may be terminated"
    fi ;;
  *) sed -n 2,7p "$0"; exit 2 ;;
esac