#!/usr/bin/env bash
set -euo pipefail

# Migrate PlantAI-dev code from /data_old (HDD) to /data (SSD), while keeping
# input/workspace on /data_old. This script is intentionally conservative:
#   - default mode is dry-run
#   - old code is renamed, not deleted
#   - old code path is replaced by a compatibility symlink after migration
#
# This version does not migrate old input/workspace content. It creates a fresh
# data root at NEW_DATA and links the new code project to NEW_DATA/input and
# NEW_DATA/workspace.

APPLY=0
STOP_SERVICES=0

OLD_CODE="/data_old/home/xia-lab/project_roots/PlantAI-dev"
NEW_CODE="/data/project_roots/PlantAI-dev"
NEW_DATA="/data_old/home/xia-lab/plantai-dev-data"
HOME_LINK="/home/xia-lab/PlantAI-dev"
LEGACY_LINK="/home/xia-lab/wutx/00test/PlantAI-dev"

usage() {
  cat <<'EOF'
Usage:
  bash tools/migrate_plantai_dev_to_data_ssd.sh [--apply] [--stop-services]

Default is dry-run. Nothing is changed unless --apply is provided.

What --apply does:
  1. Optionally stops PlantAI-dev web/Celery processes when --stop-services is used.
  2. Copies code from /data_old/home/xia-lab/project_roots/PlantAI-dev to /data/project_roots/PlantAI-dev.
  3. Creates fresh input/workspace directories under /data_old/home/xia-lab/plantai-dev-data.
  4. Creates input/workspace symlinks inside the new code directory.
  5. Renames the old code directory to PlantAI-dev.migrated_<timestamp>.
  6. Creates compatibility symlinks at old/common entry paths pointing to the new code directory.

It never deletes the old renamed project.
EOF
}

for arg in "$@"; do
  case "$arg" in
    --apply)
      APPLY=1
      ;;
    --stop-services)
      STOP_SERVICES=1
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "[ERROR] Unknown argument: $arg" >&2
      usage
      exit 2
      ;;
  esac
done

run() {
  if [[ "$APPLY" -eq 1 ]]; then
    echo "+ $*"
    "$@"
  else
    echo "[dry-run] $*"
  fi
}

require_path() {
  local path="$1"
  local label="$2"
  if [[ ! -e "$path" ]]; then
    echo "[ERROR] Missing $label: $path" >&2
    exit 1
  fi
}

echo "=== PlantAI-dev migration plan ==="
echo "OLD_CODE     = $OLD_CODE"
echo "NEW_CODE     = $NEW_CODE"
echo "NEW_DATA     = $NEW_DATA"
echo "HOME_LINK    = $HOME_LINK"
echo "LEGACY_LINK  = $LEGACY_LINK"
echo "APPLY        = $APPLY"
echo "STOP_SERVICES= $STOP_SERVICES"
echo

require_path "$OLD_CODE" "old code directory"

if [[ ! -d "$OLD_CODE/apps" || ! -d "$OLD_CODE/gfa_core" || ! -d "$OLD_CODE/scripts" ]]; then
  echo "[ERROR] OLD_CODE does not look like the full PlantAI-dev code project: $OLD_CODE" >&2
  exit 1
fi

echo "=== Disk summary ==="
df -hT /data /data_old / || true
echo

echo "=== Current sizes ==="
du -sh "$OLD_CODE" 2>/dev/null || true
du -sh "$NEW_DATA/input" "$NEW_DATA/workspace" 2>/dev/null || true
echo

if [[ "$STOP_SERVICES" -eq 1 ]]; then
  echo "=== Stop PlantAI-dev services ==="
  run pkill -f "apps.web.vue3.main:app" || true
  run pkill -f "apps/web/vue3/main.py" || true
  run pkill -f "celery.*plantai_dev" || true
  run pkill -f "celery_app" || true
  if [[ "$APPLY" -eq 1 ]]; then
    sleep 2
  fi
fi

echo "=== Process check ==="
if ps -ef | grep -E "apps.web.vue3.main|apps/web/vue3/main.py|celery.*plantai_dev|celery_app" | grep -v grep; then
  if [[ "$STOP_SERVICES" -eq 0 ]]; then
    echo "[ERROR] PlantAI-dev-related processes are still running." >&2
    echo "        Re-run with --stop-services, or stop them manually before --apply." >&2
    if [[ "$APPLY" -eq 1 ]]; then
      exit 1
    fi
  fi
else
  echo "No matching PlantAI-dev web/Celery processes found."
fi
echo

echo "=== Prepare directories ==="
if [[ "$APPLY" -eq 1 ]]; then
  new_code_parent="$(dirname "$NEW_CODE")"
  if [[ ! -w "$new_code_parent" ]]; then
    echo "[ERROR] $new_code_parent is not writable by $(id -un). Create it and chown it first, for example:" >&2
    echo "        sudo mkdir -p /data/project_roots" >&2
    echo "        sudo chown -R $(id -un):$(id -gn) /data/project_roots" >&2
    exit 1
  fi
fi
run mkdir -p "$(dirname "$NEW_CODE")"
run mkdir -p "$NEW_DATA/input" "$NEW_DATA/workspace"

if [[ -e "$NEW_CODE" && ! -d "$NEW_CODE/apps" ]]; then
  echo "[ERROR] NEW_CODE exists but does not look like a PlantAI-dev project: $NEW_CODE" >&2
  exit 1
fi

echo "=== Copy code to SSD ==="
run rsync -aHAX --info=progress2 \
  --exclude='/input' \
  --exclude='/workspace' \
  --exclude='/workspace.local_created_during_relink_*' \
  "$OLD_CODE"/ "$NEW_CODE"/

echo "=== Data directories ==="
echo "Using data root: $NEW_DATA"
echo "The script will not migrate old input/workspace content."

echo "=== Link new code to HDD data ==="
run ln -sfn "$NEW_DATA/input" "$NEW_CODE/input"
run ln -sfn "$NEW_DATA/workspace" "$NEW_CODE/workspace"

STAMP="$(date +%Y%m%d_%H%M%S)"
RENAMED_OLD="${OLD_CODE}.migrated_${STAMP}"

echo "=== Rename old code project, keep compatibility link ==="
if [[ "$APPLY" -eq 1 ]]; then
  if [[ -L "$OLD_CODE" ]]; then
    echo "[ERROR] OLD_CODE is already a symlink, refusing to rename: $OLD_CODE" >&2
    exit 1
  fi
fi
run mv "$OLD_CODE" "$RENAMED_OLD"
run ln -sfn "$NEW_CODE" "$OLD_CODE"

echo "=== Update user-facing links ==="
run ln -sfn "$NEW_CODE" "$HOME_LINK"
run ln -sfn "$NEW_CODE" "$LEGACY_LINK"

echo "=== Final verification ==="
if [[ "$APPLY" -eq 1 ]]; then
  echo "New code real path:"
  readlink -f "$NEW_CODE"
  echo "Data links:"
  readlink -f "$NEW_CODE/input"
  readlink -f "$NEW_CODE/workspace"
  echo "Old code compatibility link:"
  ls -ld "$OLD_CODE"
  echo "Renamed old code kept at:"
  ls -ld "$RENAMED_OLD"
  test -d "$NEW_CODE/apps"
  test -d "$NEW_CODE/gfa_core"
  test -d "$NEW_CODE/scripts"
  test -d "$NEW_CODE/input"
  test -d "$NEW_CODE/workspace"
  echo "[OK] Migration completed. Start services from: $NEW_CODE"
else
  echo "[dry-run] No changes made. Re-run with --apply to perform migration."
fi