social-data-attribution-w2/scripts/bootstrap/setup_worktree.sh

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

# Prepare a git worktree for running jobs.
# Symlinks shared local-path dependencies (bergson, olmes) from the main
# checkout and creates a venv via uv.
#
# Usage:
#   cd ~/dev/data-attribution-soc90
#   bash scripts/bootstrap/setup_worktree.sh
#   uv run --no-sync pytest -q
#
# On HPC (PACE ICE), uv manages its own Python 3.12 install.
# Locally, uv uses whatever Python is on PATH.

WORKTREE_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
COMMON_DIR="$(git rev-parse --git-common-dir)"
MAIN_REPO="$(cd "$COMMON_DIR/.." && pwd)"

echo "Worktree root: $WORKTREE_ROOT"
echo "Main repo:     $MAIN_REPO"

if [ "$WORKTREE_ROOT" = "$MAIN_REPO" ]; then
    echo "Error: run this from a worktree, not the main checkout." >&2
    exit 1
fi

BOOTSTRAP_SCRIPT="$WORKTREE_ROOT/scripts/bootstrap_local_deps.sh"
if [ ! -f "$BOOTSTRAP_SCRIPT" ]; then
    echo "Error: bootstrap script not found: $BOOTSTRAP_SCRIPT" >&2
    exit 1
fi

DEPENDENCY_ROOT="$MAIN_REPO" PATCH_ROOT="$WORKTREE_ROOT" bash "$BOOTSTRAP_SCRIPT"

link_dep() {
    local name="$1"
    local target="$MAIN_REPO/$name"

    if [ -L "$WORKTREE_ROOT/$name" ] && [ ! -e "$WORKTREE_ROOT/$name" ]; then
        echo "Removing broken symlink: $name"
        unlink "$WORKTREE_ROOT/$name"
    fi

    if [ -e "$WORKTREE_ROOT/$name" ]; then
        echo "Already exists: $name (skipping)"
        return
    fi

    if [ ! -d "$target" ]; then
        echo "Error: dependency source missing after bootstrap: $target" >&2
        exit 1
    fi

    ln -s "$target" "$WORKTREE_ROOT/$name"
    echo "Linked: $name -> $target"
}

link_dep bergson
link_dep olmes

copy_if_changed() {
    local relpath="$1"
    local src="$MAIN_REPO/$relpath"
    local dst="$WORKTREE_ROOT/$relpath"

    if [ ! -f "$src" ]; then
        return
    fi
    if [ -f "$dst" ] && cmp -s "$src" "$dst"; then
        echo "Up to date: $relpath (skipping)"
        return
    fi

    mkdir -p "$(dirname "$dst")"
    /bin/cp "$src" "$dst"
    echo "Copied: $relpath"
}

copy_if_changed .claude/settings.local.json
copy_if_changed .claude/autopilot/state.json
copy_if_changed uv.lock
if [ ! -f "$WORKTREE_ROOT/uv.lock" ]; then
    echo "Error: uv.lock not found in main repo or worktree. Run 'uv lock' in the main repo first." >&2
    exit 1
fi

if ! command -v uv >/dev/null 2>&1; then
    echo "Error: uv not found. Install it first." >&2
    exit 1
fi

compute_fingerprint() {
    if command -v shasum >/dev/null 2>&1; then
        cat "$WORKTREE_ROOT/pyproject.toml" "$WORKTREE_ROOT/uv.lock" | shasum -a 256 | cut -d' ' -f1
    else
        cat "$WORKTREE_ROOT/pyproject.toml" "$WORKTREE_ROOT/uv.lock" | sha256sum | cut -d' ' -f1
    fi
}

SCRATCH_DIR="$HOME/scratch"
HAS_SCRATCH=false
if [ "$(uname -s)" != "Darwin" ] && [ -d "$SCRATCH_DIR" ]; then
    HAS_SCRATCH=true
fi

if [ -z "${UV_CACHE_DIR:-}" ]; then
    if $HAS_SCRATCH; then
        export UV_CACHE_DIR="$SCRATCH_DIR/uv-cache"
    elif [ "$(uname -s)" = "Linux" ]; then
        export UV_CACHE_DIR="${TMPDIR:-/tmp}/uv-cache"
    fi
    if [ -n "${UV_CACHE_DIR:-}" ]; then
        mkdir -p "$UV_CACHE_DIR"
    fi
fi

echo "Creating venv and syncing dependencies..."
cd "$WORKTREE_ROOT"
if [ ! -d "$WORKTREE_ROOT/.venv" ]; then
    if $HAS_SCRATCH; then
        WORKTREE_NAME="$(basename "$WORKTREE_ROOT")"
        VENV_DIR="$SCRATCH_DIR/venvs/$WORKTREE_NAME"
        mkdir -p "$(dirname "$VENV_DIR")"
        uv venv "$VENV_DIR"
        ln -s "$VENV_DIR" "$WORKTREE_ROOT/.venv"
        echo "Venv created in scratch: $VENV_DIR (symlinked to .venv)"
    else
        uv venv
    fi
fi

FINGERPRINT_FILE="$WORKTREE_ROOT/.venv/.bootstrap_fingerprint"
current_fp="$(compute_fingerprint)"

if [ -f "$FINGERPRINT_FILE" ] && [ -x "$WORKTREE_ROOT/.venv/bin/python" ] && [ "$(cat "$FINGERPRINT_FILE")" = "$current_fp" ]; then
    echo "Dependencies unchanged (fingerprint match). Skipping uv sync."
else
    # TODO: add torchopt to pyproject.toml or lockfile so this is automatic.
    # bergson imports torchopt at CLI entry (bergson.magic) but it's not in
    # the frozen lockfile. Without it, bergson preconditioners/build commands fail.
    uv pip install setuptools pybind11 torchopt
    if ! uv sync --frozen --no-build-isolation; then
        echo "uv sync failed, retrying without cache..."
        UV_NO_CACHE=1 uv sync --frozen --no-build-isolation
    fi
    echo "$current_fp" > "$FINGERPRINT_FILE"
fi
echo "Done. Worktree ready at $WORKTREE_ROOT"
echo "Use 'uv run --no-sync ...' for routine verification commands in this worktree."