uw

#!/bin/bash
#
# Underworld3 environment wrapper
#
set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONFIG_FILE="$SCRIPT_DIR/.pixi-env"
PETSC_CUSTOM="$SCRIPT_DIR/petsc-custom/petsc"

# Ensure pixi resolves to THIS repo/worktree, not a parent env
export PIXI_PROJECT_MANIFEST="$SCRIPT_DIR/pixi.toml"

# Colors (if terminal supports it)
if [ -t 1 ]; then
    GREEN='\033[0;32m'
    YELLOW='\033[1;33m'
    BOLD='\033[1m'
    NC='\033[0m'
else
    GREEN='' YELLOW='' BOLD='' NC=''
fi

# Find pixi - check PATH first, then ~/.pixi/bin/ (where installer puts it)
find_pixi() {
    if command -v pixi &> /dev/null; then
        echo "pixi"
    elif [ -x "$HOME/.pixi/bin/pixi" ]; then
        echo "$HOME/.pixi/bin/pixi"
    else
        echo ""
    fi
}

PIXI=$(find_pixi)

get_env() {
    if [ -f "$CONFIG_FILE" ]; then
        cat "$CONFIG_FILE" | tr -d '[:space:]'
    else
        echo "runtime"
    fi
}

petsc_version_short() {
    # Read active PETSc version from .petsc-version, or detect from checkout.
    local ver_file="$SCRIPT_DIR/petsc-custom/.petsc-version"
    if [ -f "$ver_file" ]; then
        cat "$ver_file" | tr -d '[:space:]'
        return
    fi
    # Detect from git tag
    if [ -d "$PETSC_CUSTOM/.git" ]; then
        local tag
        tag=$(cd "$PETSC_CUSTOM" && git describe --tags --abbrev=0 2>/dev/null || echo "")
        if [ -n "$tag" ]; then
            echo "$tag" | sed -E 's/^v([0-9]+)\.([0-9]+).*/\1\2/'
            return
        fi
    fi
    echo "324"  # default
}

petsc_arch_for_env() {
    # Derive the PETSC_ARCH from the environment name and active PETSc version.
    # Format: petsc-{version}-uw-{mpi}  (e.g. petsc-324-uw-openmpi)
    local env="$1"
    local ver=$(petsc_version_short)
    local mpi=""

    case "$env" in
        amr-mpich*)   mpi="mpich" ;;
        amr-openmpi*) mpi="openmpi" ;;
        amr*)
            # Platform-default MPI
            if [[ "$(uname -s)" == "Darwin" ]]; then
                mpi="openmpi"
            else
                mpi="mpich"
            fi
            ;;
        *)            echo ""; return ;;  # conda-forge, no custom PETSc
    esac

    # Check for versioned arch first, fall back to legacy unversioned
    local versioned="petsc-${ver}-uw-${mpi}"
    local legacy="petsc-4-uw-${mpi}"
    if [ -d "$PETSC_CUSTOM/${versioned}" ]; then
        echo "${versioned}"
    elif [ -d "$PETSC_CUSTOM/${legacy}" ]; then
        echo "${legacy}"
    else
        # No build exists yet — return the versioned name for new builds
        echo "${versioned}"
    fi
}

petsc_built() {
    # Check for actual PETSc library, not just the lib directory
    # The lib directory can exist if the build failed partway through
    local env="${1:-$(get_env)}"
    local arch=$(petsc_arch_for_env "$env")
    [ -z "$arch" ] && return 1  # not an AMR env
    [ -f "$PETSC_CUSTOM/$arch/lib/libpetsc.dylib" ] || \
    [ -f "$PETSC_CUSTOM/$arch/lib/libpetsc.so" ] || \
    [ -f "$PETSC_CUSTOM/$arch/lib/libpetsc.a" ]
}

petsc4py_installed() {
    local env="$1"
    $PIXI run -e "$env" python -c "import petsc4py" 2>/dev/null
}

uw3_importable() {
    local env="$1"
    $PIXI run -e "$env" python -c "import underworld3" 2>/dev/null
}

is_amr_env() {
    local env="$1"
    [[ "$env" == amr* ]]
}

# Smart build with dependency chain handling
run_build() {
    local env=$(get_env)

    echo -e "${BOLD}Building underworld3${NC} (environment: $env)"
    echo ""

    # AMR environments need custom PETSc + petsc4py first
    if is_amr_env "$env"; then
        # Step 1: Check/build custom PETSc
        if ! petsc_built; then
            echo -e "${YELLOW}AMR environment requires custom PETSc (not yet built)${NC}"
            echo ""
            echo "Custom PETSc includes adaptive mesh tools (pragmatic, mmg, parmmg)"
            echo "Build time: ~1 hour on Apple Silicon"
            echo ""
            read -p "Build custom PETSc now? [Y/n]: " build_petsc
            build_petsc=${build_petsc:-y}
            if [ "$build_petsc" = "y" ] || [ "$build_petsc" = "Y" ]; then
                echo ""
                $PIXI run -e "$env" ./petsc-custom/build-petsc.sh
                if [ $? -ne 0 ]; then
                    echo -e "${YELLOW}PETSc build failed. Check errors above.${NC}"
                    exit 1
                fi
            else
                echo ""
                echo "Cannot build underworld3 without PETSc."
                echo "Run './uw petsc-local-build' when ready, then './uw build'"
                exit 1
            fi
        else
            echo -e "  ${GREEN}✓${NC} Custom PETSc already built"
        fi

        # Step 2: Check/build petsc4py
        if ! petsc4py_installed "$env"; then
            echo "  Installing petsc4py for $env..."
            # The in-tree build/ cache is SHARED across envs/worktrees and pip
            # silently reuses its objects even when numpy/PETSc changed —
            # stale-ABI imports ("numpy.dtype size changed"). Always start clean.
            rm -rf "$PETSC_CUSTOM/src/binding/petsc4py/build"
            $PIXI run -e "$env" pip install "$PETSC_CUSTOM/src/binding/petsc4py" --no-build-isolation --no-cache-dir || {
                echo -e "${YELLOW}petsc4py build failed${NC}"
                exit 1
            }
            echo -e "  ${GREEN}✓${NC} petsc4py installed"
        else
            echo -e "  ${GREEN}✓${NC} petsc4py already installed"
        fi
    else
        # Non-AMR: conda-forge petsc4py should be there from pixi install
        if ! petsc4py_installed "$env"; then
            echo -e "${YELLOW}petsc4py not found. Running pixi install...${NC}"
            $PIXI install -e "$env"
        fi
    fi

    # Step 3: Build underworld3
    # Clean build cache if PETSc target changed (prevents stale .so from
    # being installed when switching between default and AMR environments)
    local petsc_marker="$SCRIPT_DIR/build/.petsc_target"
    local current_target
    current_target=$($PIXI run -e "$env" python -c "import petsc4py; c=petsc4py.get_config(); print(c['PETSC_DIR']+'/'+c.get('PETSC_ARCH',''))" 2>/dev/null)
    local need_clean=false
    if [ -f "$petsc_marker" ]; then
        local prev_target
        prev_target=$(cat "$petsc_marker")
        if [ "$current_target" != "$prev_target" ]; then
            need_clean=true
        fi
    elif ls "$SCRIPT_DIR"/build/lib.* 1>/dev/null 2>&1; then
        # No marker but build cache exists — stale from a raw pip install
        need_clean=true
    fi
    if [ "$need_clean" = true ]; then
        echo "  PETSc target changed — cleaning build cache..."
        rm -rf "$SCRIPT_DIR"/build/lib.* "$SCRIPT_DIR"/build/temp.*
    fi
    mkdir -p "$SCRIPT_DIR/build"
    echo "$current_target" > "$petsc_marker"

    echo "  Building underworld3..."
    # Determine the ACTUAL site-packages that Python will search.
    # pip may resolve a different path (e.g. through worktree symlinks),
    # so we ask Python directly and install with --target.
    local site_packages
    site_packages=$($PIXI run -e "$env" python -c "import site; print(site.getsitepackages()[0])" 2>/dev/null)
    if [ -z "$site_packages" ]; then
        echo -e "${YELLOW}Could not determine site-packages path${NC}"
        exit 1
    fi

    # Remove any stale editable install markers (pip install -e leaves .pth files
    # that redirect imports to the source tree, breaking environment isolation)
    find "$site_packages" -maxdepth 1 -name "__editable__*underworld*" -delete 2>/dev/null

    # Remove any stale .so from the source tree (left by editable installs)
    find "$SCRIPT_DIR/src/underworld3" -name "*.so" -delete 2>/dev/null

    # --no-cache-dir: UW3 is always version 0.0.0, so pip's wheel cache
    # treats every build as "already cached" and silently reuses stale code.
    # --target: install to where Python ACTUALLY looks, not where pip thinks.
    # NEVER use -e (editable) — it breaks worktree and environment isolation.
    $PIXI run -e "$env" pip install . --no-build-isolation --no-cache-dir \
        --upgrade --target="$site_packages" || {
        echo -e "${YELLOW}underworld3 build failed${NC}"
        exit 1
    }

    # Verify the install is importable
    $PIXI run -e "$env" python -c "import underworld3" 2>/dev/null || {
        echo -e "${YELLOW}underworld3 installed but cannot be imported!${NC}"
        echo "  site-packages: $site_packages"
        echo "  Check for stale .pth files or .so mismatches."
        exit 1
    }

    # Verify Python source files were actually installed.
    # pip's wheel cache can silently serve stale .py files even with
    # --no-cache-dir (version 0.0.0 problem). Compare checksums of key
    # source files against the installed copies and warn if they differ.
    local site_packages=$($PIXI run -e "$env" python -c \
        "import underworld3, pathlib; print(pathlib.Path(underworld3.__file__).parent)" 2>/dev/null)

    if [ -n "$site_packages" ] && [ -d "$site_packages" ]; then
        local stale_files=()
        # Check .py files under src/underworld3/ against installed copies
        while IFS= read -r src_file; do
            local rel_path="${src_file#src/underworld3/}"
            local inst_file="$site_packages/$rel_path"
            if [ -f "$inst_file" ]; then
                local src_hash=$(shasum -a 256 "$src_file" | cut -d' ' -f1)
                local inst_hash=$(shasum -a 256 "$inst_file" | cut -d' ' -f1)
                if [ "$src_hash" != "$inst_hash" ]; then
                    stale_files+=("$rel_path")
                fi
            fi
        done < <(find src/underworld3 -name '*.py' -not -path '*/__pycache__/*' 2>/dev/null)

        if [ ${#stale_files[@]} -gt 0 ]; then
            echo ""
            echo -e "${YELLOW}${BOLD}WARNING: ${#stale_files[@]} source file(s) differ from installed copies:${NC}"
            for f in "${stale_files[@]}"; do
                echo -e "  ${YELLOW}STALE: $f${NC}"
            done
            echo ""
            echo "  Fixing: copying source files to installed location..."
            for f in "${stale_files[@]}"; do
                cp "src/underworld3/$f" "$site_packages/$f"
            done
            echo -e "  ${GREEN}✓${NC} ${#stale_files[@]} file(s) updated"
        fi
    fi

    # JIT cache env-fingerprint check.
    #
    # The on-disk JIT cache (~/.cache/underworld3/jit/) holds compiled .so
    # files keyed by C-source hash + (PETSc, UW3) version salt. Anything
    # outside that salt — compiler identity, CFLAGS, Python ABI, pixi env —
    # could silently invalidate cached binaries (e.g. switching pixi env
    # from default to amr-dev links against different libc++ shims).
    #
    # We detect such changes by writing an env fingerprint to the cache dir
    # and wiping the cache on mismatch.
    local jit_cache_dir="${UW_JIT_CACHE_DIR:-$HOME/.cache/underworld3/jit}"
    local fp_file="$jit_cache_dir/.env-fingerprint"
    local cc_id python_id pixi_env_id
    cc_id=$(${CC:-cc} --version 2>/dev/null | head -1)
    python_id=$($PIXI run -e "$env" python -c "import sys; print(sys.implementation.name, sys.version)" 2>/dev/null)
    pixi_env_id="$env"
    local fp
    fp=$(printf '%s\nCC=%s\nCFLAGS=%s\nPYTHON=%s\nPIXI_ENV=%s\n' \
            "fingerprint-v1" "$cc_id" "${CFLAGS:-}" "$python_id" "$pixi_env_id" \
            | shasum -a 256 | cut -c1-32)
    if [ -f "$fp_file" ]; then
        local prev_fp
        prev_fp=$(cat "$fp_file" 2>/dev/null)
        if [ "$prev_fp" != "$fp" ]; then
            echo "  Compiler / Python / pixi env changed; wiping JIT cache..."
            find "$jit_cache_dir" -maxdepth 1 \
                \( -name '*.so' -o -name '*.manifest.json' -o -name '.*.lock' \
                   -o -name '.inprogress_*' \) -delete 2>/dev/null
        fi
    fi
    mkdir -p "$jit_cache_dir"
    echo "$fp" > "$fp_file"

    echo ""
    echo -e "${GREEN}${BOLD}Build complete!${NC}"
    echo ""
    echo "Next: ./uw test    (run tests)"
    echo "      ./uw doctor  (check configuration)"
}

run_petsc_cmd() {
    local petsc_cmd="${1:-versions}"
    local arg="$2"
    case "$petsc_cmd" in
        versions|list)
            $PIXI run -e "$(get_env)" "$SCRIPT_DIR/petsc-custom/build-petsc.sh" versions
            ;;
        switch)
            if [ -z "$arg" ]; then
                echo "Usage: ./uw petsc switch <version-tag>"
                echo "  e.g.: ./uw petsc switch v3.25.0"
                exit 1
            fi
            $PIXI run -e "$(get_env)" "$SCRIPT_DIR/petsc-custom/build-petsc.sh" checkout "$arg"

            # Nuke stale build artifacts and petsc4py. The next ./uw build's
            # target-change detection can't fire because the still-installed
            # petsc4py reports the OLD arch, so we force a full rebuild here.
            # (run_build already passes --no-cache-dir to pip install . so we
            # don't need a global pip cache purge — only petsc4py's local-path
            # install needs its own --no-cache-dir.)
            echo ""
            echo "Clearing stale build artifacts for clean rebuild..."
            rm -rf "$SCRIPT_DIR"/build/lib.* "$SCRIPT_DIR"/build/temp.* "$SCRIPT_DIR"/build/bdist.* 2>/dev/null
            rm -f "$SCRIPT_DIR/build/.petsc_target"
            if is_amr_env "$(get_env)"; then
                echo "Uninstalling old petsc4py (linked to previous PETSc build)..."
                $PIXI run -e "$(get_env)" pip uninstall -y petsc4py 2>/dev/null || true
            fi

            echo ""
            echo "To complete the switch, rebuild PETSc then underworld3:"
            echo "  ./uw petsc build"
            echo "  ./uw build"
            ;;
        active)
            local ver=$(petsc_version_short)
            local arch=$(petsc_arch_for_env "$(get_env)")
            echo "Active PETSc: ${ver} (arch: ${arch})"
            ;;
        build)
            $PIXI run -e "$(get_env)" "$SCRIPT_DIR/petsc-custom/build-petsc.sh"
            ;;
        *)
            echo "Usage: ./uw petsc [versions|switch <tag>|active|build]"
            echo ""
            echo "  versions     List available PETSc builds"
            echo "  switch <tag> Checkout a PETSc version (e.g. v3.25.0)"
            echo "  active       Show active PETSc version and arch"
            echo "  build        Build PETSc for the active version"
            ;;
    esac
}

# Standalone diagnostics (works even if underworld3 not built)
run_doctor() {
    local env=$(get_env)
    local errors=0
    local warnings=0

    echo ""
    echo "============================================================"
    echo "  Underworld3 Diagnostics"
    echo "============================================================"
    echo ""
    echo "Environment: $env"
    echo ""

    # Check 1: Pixi environment installed?
    echo -n "  Checking pixi environment... "
    if env_installed "$env"; then
        echo -e "${GREEN}[OK]${NC} $env installed"
    else
        echo -e "${YELLOW}[ERROR]${NC} $env not installed"
        echo "      Fix: pixi install -e $env"
        ((errors++))
    fi

    # Check 2: For AMR, is custom PETSc built?
    if is_amr_env "$env"; then
        echo -n "  Checking custom PETSc... "
        if petsc_built; then
            echo -e "${GREEN}[OK]${NC} Custom PETSc built"
        else
            echo -e "${YELLOW}[ERROR]${NC} Custom PETSc not built"
            echo "      Fix: ./uw petsc-local-build  (takes ~1 hour)"
            ((errors++))
        fi

        # Check 3: petsc4py installed?
        echo -n "  Checking petsc4py... "
        if petsc4py_installed "$env"; then
            echo -e "${GREEN}[OK]${NC} petsc4py available"
        else
            echo -e "${YELLOW}[ERROR]${NC} petsc4py not installed"
            if petsc_built; then
                echo "      Fix: ./uw build  (will install petsc4py)"
            else
                echo "      Fix: Build PETSc first, then ./uw build"
            fi
            ((errors++))
        fi
    else
        # Non-AMR: check conda-forge petsc4py
        echo -n "  Checking petsc4py... "
        if petsc4py_installed "$env"; then
            echo -e "${GREEN}[OK]${NC} petsc4py available"
        else
            echo -e "${YELLOW}[ERROR]${NC} petsc4py not found"
            echo "      Fix: pixi install -e $env"
            ((errors++))
        fi
    fi

    # Check 4: underworld3 importable?
    echo -n "  Checking underworld3... "
    local import_error=""
    local import_status=0
    import_error=$($PIXI run -e "$env" python -c "import underworld3" 2>&1) || import_status=$?

    if [ $import_status -eq 0 ]; then
        echo -e "${GREEN}[OK]${NC} underworld3 importable"

        # If importable, run Python diagnostics for deeper checks
        echo ""
        echo "  Running detailed diagnostics..."
        $PIXI run -e "$env" python -c "import underworld3 as uw; uw.doctor()" 2>&1 | sed 's/^/  /'
    else
        echo -e "${YELLOW}[ERROR]${NC} underworld3 not importable"

        # Check for specific PETSc library mismatch error
        if echo "$import_error" | grep -q "libpetsc"; then
            echo ""
            echo -e "      ${YELLOW}PETSc library version mismatch detected!${NC}"
            echo ""
            # Extract the library name from error
            local lib_name=$(echo "$import_error" | grep -o "libpetsc[^']*" | head -1)
            if [ -n "$lib_name" ]; then
                echo "      Extension was built against: $lib_name"
            fi
            # Get current PETSc version
            local petsc_ver=$($PIXI run -e "$env" python -c "from petsc4py import PETSc; v=PETSc.Sys.getVersion(); print(f'{v[0]}.{v[1]}.{v[2]}')" 2>/dev/null)
            if [ -n "$petsc_ver" ]; then
                echo "      Current petsc4py version:    $petsc_ver"
            fi
            echo ""
            echo "      This happens when PETSc is updated but underworld3"
            echo "      wasn't rebuilt. The extension is linked to the old library."
            echo ""
            echo -e "      ${BOLD}Fix:${NC}"
            echo "        pip uninstall underworld3"
            echo "        rm -rf build/"
            echo "        ./uw build"
            echo ""
        else
            echo "      Fix: ./uw build"
            # Show first few lines of error for debugging
            echo ""
            echo "      Import error:"
            echo "$import_error" | head -5 | sed 's/^/        /'
        fi
        ((errors++))
    fi

    echo ""
    echo "------------------------------------------------------------"
    if [ $errors -gt 0 ]; then
        echo -e "  ${YELLOW}[FAIL]${NC} $errors error(s) found - see fixes above"
        echo ""
        echo "  Quick fix: ./uw build  (handles dependency chain)"
    else
        echo -e "  ${GREEN}[OK]${NC} All checks passed"
    fi
    echo "------------------------------------------------------------"
}

# Check and install pixi if needed
check_pixi() {
    if [ -n "$PIXI" ]; then
        return 0
    fi

    echo -e "${YELLOW}Pixi package manager not found.${NC}"
    echo ""
    echo "Pixi is required to manage Underworld3 dependencies."
    echo "Install from: https://pixi.sh"
    echo ""
    read -p "Install pixi now? [Y/n]: " install_pixi
    install_pixi=${install_pixi:-y}

    if [ "$install_pixi" = "y" ] || [ "$install_pixi" = "Y" ]; then
        echo ""
        echo "Installing pixi..."
        curl -fsSL https://pixi.sh/install.sh | bash

        # Update PIXI variable to find newly installed pixi
        PIXI=$(find_pixi)

        if [ -n "$PIXI" ]; then
            echo -e "${GREEN}✓ Pixi installed successfully${NC}"
            echo ""
            return 0
        else
            echo ""
            echo -e "${YELLOW}Pixi installation failed.${NC}"
            echo "Please install manually from https://pixi.sh"
            exit 1
        fi
    else
        echo ""
        echo "Install pixi manually from https://pixi.sh then run './uw setup'"
        exit 1
    fi
}

# Configure AI tools path (standalone)
configure_ai_tools() {
    echo -e "${BOLD}Extended AI Context Configuration${NC}"
    echo ""
    echo "Configure paths to additional AI instruction files (.md) that live"
    echo "outside this repository. Useful for personal coding preferences,"
    echo "project planning systems, or private instructions."
    echo ""
    
    # Show current setting if any
    if [ -n "$UW_AI_TOOLS_PATH" ]; then
        echo "Current setting: $UW_AI_TOOLS_PATH"
        echo ""
    fi
    
    echo "Enter path(s) to directories containing .md instruction files."
    echo "Separate multiple paths with colons (like PATH)."
    echo "Leave empty to remove/skip."
    echo ""
    echo "Example: ~/Library/CloudStorage/Box-Box/Planning:~/my-ai-notes"
    echo ""
    read -p "Path(s): " ai_tools_path
    
    # Detect shell config file
    local shell_config=""
    if [ -f "$HOME/.zshrc" ]; then
        shell_config="$HOME/.zshrc"
    elif [ -f "$HOME/.bashrc" ]; then
        shell_config="$HOME/.bashrc"
    elif [ -f "$HOME/.bash_profile" ]; then
        shell_config="$HOME/.bash_profile"
    fi
    
    if [ -z "$ai_tools_path" ]; then
        # Empty input - offer to remove existing
        if [ -n "$shell_config" ] && grep -q "^export UW_AI_TOOLS_PATH=" "$shell_config" 2>/dev/null; then
            echo ""
            read -p "Remove existing UW_AI_TOOLS_PATH from $shell_config? [y/N]: " remove_it
            if [ "$remove_it" = "y" ] || [ "$remove_it" = "Y" ]; then
                grep -v "^export UW_AI_TOOLS_PATH=" "$shell_config" | grep -v "^# Underworld3 AI tools path" > "$shell_config.tmp"
                mv "$shell_config.tmp" "$shell_config"
                echo -e "${GREEN}✓${NC} Removed from $shell_config"
            fi
        else
            echo "No path configured."
        fi
        return
    fi
    
    # Expand ~ to $HOME
    ai_tools_path="${ai_tools_path//\~/$HOME}"
    
    if [ -n "$shell_config" ]; then
        echo ""
        read -p "Add to $shell_config? [Y/n]: " add_to_shell
        add_to_shell=${add_to_shell:-y}
        
        if [ "$add_to_shell" = "y" ] || [ "$add_to_shell" = "Y" ]; then
            # Remove any existing UW_AI_TOOLS_PATH lines
            grep -v "^export UW_AI_TOOLS_PATH=" "$shell_config" | grep -v "^# Underworld3 AI tools path" > "$shell_config.tmp" 2>/dev/null || true
            mv "$shell_config.tmp" "$shell_config"
            
            # Add new line
            echo "" >> "$shell_config"
            echo "# Underworld3 AI tools path (configured by ./uw ai-tools)" >> "$shell_config"
            echo "export UW_AI_TOOLS_PATH=\"$ai_tools_path\"" >> "$shell_config"
            
            echo -e "${GREEN}✓${NC} Added to $shell_config"
            echo ""
            echo "Run 'source $shell_config' or start a new terminal to activate."
        else
            echo ""
            echo "Add this to your shell config manually:"
            echo "  export UW_AI_TOOLS_PATH=\"$ai_tools_path\""
        fi
    else
        echo ""
        echo "Could not detect shell config. Add this manually:"
        echo "  export UW_AI_TOOLS_PATH=\"$ai_tools_path\""
    fi
}

# Configure Claude Code permissions (safe defaults for development monitoring)
# Usage: configure_claude_permissions <pixi-env>
configure_claude_permissions() {
    local env="${1:-$(get_env)}"
    mkdir -p "$HOME/.claude"

    # Use Python to safely merge permissions into settings.json
    $PIXI run -e "$env" python -c "
import json, os, sys

settings_file = os.path.expanduser('~/.claude/settings.json')

# Load existing settings or start fresh
settings = {}
if os.path.exists(settings_file):
    with open(settings_file) as f:
        settings = json.load(f)

# Ensure permissions.allow exists
settings.setdefault('permissions', {})
settings['permissions'].setdefault('allow', [])
allow = settings['permissions']['allow']

# Permissions for development and unattended monitoring
required = [
    'Bash(git *)',
    'Bash(gh *)',
    'Bash(pixi *)',
    'Bash(./uw *)',
    'Bash(python *)',
    'Bash(sleep *)',
    \"Bash(osascript -e 'display notification*')\",
    'Bash(notify-send *)',
    'Edit',
    'Write',
    'Agent',
]

added = []
for perm in required:
    if perm not in allow:
        allow.append(perm)
        added.append(perm)

with open(settings_file, 'w') as f:
    json.dump(settings, f, indent=2)
    print('', file=f)

if added:
    print(f'Added {len(added)} permission(s):')
    for p in added:
        print(f'  + {p}')
else:
    print('All permissions already configured.')
"
    echo ""
    echo -e "  ${GREEN}✓${NC} Claude Code permissions updated"
    echo "  Includes: git, gh, pixi, build tools, desktop notifications (macOS + Linux)"
}

# Interactive setup wizard
run_setup() {
    # Ensure pixi is available
    check_pixi

    echo -e "${BOLD}Underworld3 Setup${NC}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""

    # Question 1: AMR tools?
    echo -e "${BOLD}Do you need adaptive mesh refinement (AMR) tools?${NC}"
    echo "  (pragmatic, mmg, parmmg for anisotropic mesh adaptation)"
    echo ""
    echo "  [1] No  - conda-forge PETSc (5-minute install)"
    echo "  [2] Yes - custom PETSc build (~1 hour)"
    echo ""
    read -p "Choice [1]: " amr_choice
    amr_choice=${amr_choice:-1}

    local use_amr=false
    [ "$amr_choice" = "2" ] && use_amr=true

    echo ""

    # Question 2: Features?
    echo -e "${BOLD}What features do you need?${NC}"
    echo ""
    echo "  [1] Minimal   - build and run only"
    echo "  [2] Runtime   - add visualization + Jupyter"
    echo "  [3] Developer - add Claude, linting, docs"
    echo ""
    read -p "Choice [2]: " feature_choice
    feature_choice=${feature_choice:-2}

    echo ""

    # Question 3: MPI implementation?
    echo -e "${BOLD}Which MPI implementation?${NC}"
    echo ""
    # Platform-default MPI
    local default_mpi="mpich"
    local other_mpi="openmpi"
    if [[ "$(uname -s)" == "Darwin" ]]; then
        default_mpi="openmpi"
        other_mpi="mpich"
    fi
    echo "  [1] Platform default ($default_mpi)"
    echo "  [2] Override: $other_mpi"
    if [[ "$(uname -s)" == "Darwin" ]]; then
        echo ""
        echo "  Note: MPICH 4.x has a known scaling bug on macOS (issue #68)."
        echo "  OpenMPI is recommended unless you need MPICH for compatibility."
    fi
    echo ""
    read -p "Choice [1]: " mpi_choice
    mpi_choice=${mpi_choice:-1}

    local mpi_override=false
    [ "$mpi_choice" = "2" ] && mpi_override=true

    # Determine environment from the three choices
    # Primary envs use platform-default MPI; override envs force non-default MPI.
    local new_env
    if [ "$mpi_override" = false ]; then
        # Platform-default MPI → use primary environments
        if [ "$use_amr" = true ]; then
            case "$feature_choice" in
                1) new_env="amr" ;;
                3) new_env="amr-dev" ;;
                *) new_env="amr-runtime" ;;
            esac
        else
            case "$feature_choice" in
                1) new_env="default" ;;
                3) new_env="dev" ;;
                *) new_env="runtime" ;;
            esac
        fi
    else
        # Non-default MPI → use override environments
        if [ "$use_amr" = true ]; then
            case "$feature_choice" in
                1) new_env="amr-${other_mpi}" ;;
                3) new_env="amr-${other_mpi}-dev" ;;
                *) new_env="amr-${other_mpi}-dev" ;;  # no runtime override; use dev
            esac
        else
            case "$feature_choice" in
                1) new_env="${other_mpi}" ;;
                3) new_env="${other_mpi}-dev" ;;
                *) new_env="${other_mpi}-dev" ;;  # no runtime override; use dev
            esac
        fi
    fi

    echo ""
    echo -e "Environment: ${GREEN}$new_env${NC}"
    echo "$new_env" > "$CONFIG_FILE"
    echo ""

    # Install
    echo "Installing dependencies..."
    $PIXI install -e "$new_env"
    echo ""

    # AMR-specific
    if [ "$use_amr" = true ]; then
        if petsc_built; then
            echo -e "${GREEN}✓ Custom PETSc already built${NC}"
            echo "Installing petsc4py..."
            rm -rf "$PETSC_CUSTOM/src/binding/petsc4py/build"   # shared stale-ABI cache
            $PIXI run -e "$new_env" pip install "$PETSC_CUSTOM/src/binding/petsc4py" --no-build-isolation 2>/dev/null
        else
            echo -e "${YELLOW}Custom PETSc needs to be built (~1 hour)${NC}"
            read -p "Build now? [y/N]: " build_now
            if [ "$build_now" = "y" ] || [ "$build_now" = "Y" ]; then
                $PIXI run -e "$new_env" ./petsc-custom/build-petsc.sh
            else
                echo ""
                echo "Build later with: ./uw petsc-local-build"
                echo ""
            fi
        fi
    fi

    # Build underworld3
    echo ""
    read -p "Build underworld3? [Y/n]: " build_uw
    build_uw=${build_uw:-y}
    if [ "$build_uw" = "y" ] || [ "$build_uw" = "Y" ]; then
        run_build
    fi

    # Optional: Configure extended AI context paths
    echo ""
    echo -e "${BOLD}Extended AI Context (Optional)${NC}"
    echo "  You can configure paths to additional AI instruction files that live"
    echo "  outside this repository. This is useful for personal coding preferences,"
    echo "  project planning systems, or private instructions not shared with others."
    echo ""
    read -p "Configure AI tools path? [y/N]: " config_ai
    config_ai=${config_ai:-n}
    
    if [ "$config_ai" = "y" ] || [ "$config_ai" = "Y" ]; then
        echo ""
        echo "  Enter path(s) to directories containing .md instruction files."
        echo "  Separate multiple paths with colons (like PATH)."
        echo "  Example: ~/Library/CloudStorage/Box-Box/Planning:~/my-ai-notes"
        echo ""
        read -p "  Path(s): " ai_tools_path
        
        if [ -n "$ai_tools_path" ]; then
            # Expand ~ to $HOME
            ai_tools_path="${ai_tools_path//\~/$HOME}"
            
            # Detect shell config file
            local shell_config=""
            if [ -f "$HOME/.zshrc" ]; then
                shell_config="$HOME/.zshrc"
            elif [ -f "$HOME/.bashrc" ]; then
                shell_config="$HOME/.bashrc"
            elif [ -f "$HOME/.bash_profile" ]; then
                shell_config="$HOME/.bash_profile"
            fi
            
            if [ -n "$shell_config" ]; then
                echo ""
                read -p "  Add to $shell_config? [Y/n]: " add_to_shell
                add_to_shell=${add_to_shell:-y}
                
                if [ "$add_to_shell" = "y" ] || [ "$add_to_shell" = "Y" ]; then
                    # Remove any existing UW_AI_TOOLS_PATH line
                    grep -v "^export UW_AI_TOOLS_PATH=" "$shell_config" > "$shell_config.tmp" 2>/dev/null || true
                    mv "$shell_config.tmp" "$shell_config"
                    
                    # Add new line
                    echo "" >> "$shell_config"
                    echo "# Underworld3 AI tools path (configured by ./uw setup)" >> "$shell_config"
                    echo "export UW_AI_TOOLS_PATH=\"$ai_tools_path\"" >> "$shell_config"
                    
                    echo -e "  ${GREEN}✓${NC} Added to $shell_config"
                    echo "  Run 'source $shell_config' or start a new terminal to activate."
                else
                    echo ""
                    echo "  Add this to your shell config manually:"
                    echo "    export UW_AI_TOOLS_PATH=\"$ai_tools_path\""
                fi
            else
                echo ""
                echo "  Could not detect shell config. Add this manually:"
                echo "    export UW_AI_TOOLS_PATH=\"$ai_tools_path\""
            fi
        fi
    fi

    # Optional: Install the Claude Code CLI (only meaningful for the Developer feature,
    # which brings nodejs into the env). Reuses the install-claude pixi task.
    if [ "$feature_choice" = "3" ]; then
        echo ""
        echo -e "${BOLD}Claude Code CLI (Optional)${NC}"
        echo "  Install the Claude Code CLI globally via the dev environment's npm:"
        echo "    pixi run -e $new_env install-claude"
        echo "  (equivalent to: npm install -g @anthropic-ai/claude-code)"
        echo ""
        read -p "Install Claude Code CLI now? [y/N]: " install_claude_cli
        install_claude_cli=${install_claude_cli:-n}

        if [ "$install_claude_cli" = "y" ] || [ "$install_claude_cli" = "Y" ]; then
            $PIXI run -e "$new_env" install-claude
            echo -e "  ${GREEN}✓${NC} Claude Code CLI installed"
        else
            echo "  Skipped. Run later with: ./uw install-claude"
        fi
    fi

    # Optional: Configure Claude Code permissions for unattended monitoring
    echo ""
    echo -e "${BOLD}Claude Code Permissions (Optional)${NC}"
    echo "  Configure allowed tools for Claude Code sessions, including permissions"
    echo "  for background monitoring (CI checks, issue tracking, notifications)."
    echo "  Adds safe defaults to ~/.claude/settings.json."
    echo ""
    read -p "Configure Claude Code permissions? [y/N]: " config_perms
    config_perms=${config_perms:-n}

    if [ "$config_perms" = "y" ] || [ "$config_perms" = "Y" ]; then
        configure_claude_permissions "$new_env"
    fi

    echo ""
    echo -e "${GREEN}${BOLD}Setup complete!${NC}"
    echo ""
    echo "Usage:"
    echo "  ./uw build          Build underworld3"
    echo "  ./uw test           Run tests"
    echo "  ./uw shell          Open a shell in this env ($new_env)"

    # Show Jupyter info
    case "$new_env" in
        default|amr)
            echo ""
            echo "Note: Minimal environment - no JupyterLab/visualization."
            echo "Notebooks work via IDE kernels (VS Code, etc)."
            echo "Run './uw setup' again to add JupyterLab."
            ;;
        *)
            echo "  ./uw jupyter lab    Start JupyterLab"
            ;;
    esac
    echo ""
    echo "Note: a bare 'pixi shell' / 'pixi run' uses pixi's lightweight 'default'"
    echo "      env, which lacks dev tools (nodejs, claude, linters, docs)."
    echo "      Use './uw shell' (or 'pixi shell -e $new_env') to activate this env."
    echo ""
    echo "For VS Code / IDE Jupyter kernels, set pixi environment to: $new_env"
}

# Show comprehensive help
show_help() {
    local verbose="${1:-0}"
    cat << 'EOF'
UNDERWORLD3 ENVIRONMENT WRAPPER
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

QUICK START

  ./uw setup              Configure environment (interactive)
  ./uw build              Build underworld3
  ./uw test               Run tests (level 1, ~2 min)
  ./uw jupyter lab        Start JupyterLab
  ./uw status             Check for updates on GitHub
  ./uw update             Pull latest and rebuild

ENVIRONMENTS

  MPI is chosen automatically per platform (OpenMPI on macOS, MPICH on
  Linux). Use ./uw setup to configure — it stores the choice in .pixi-env.

  Primary environments (platform-default MPI):
  ┌─────────────┬──────────────────┬───────────────────┐
  │ Environment │ PETSc            │ Features          │
  ├─────────────┼──────────────────┼───────────────────┤
  │ default     │ conda-forge      │ minimal           │
  │ runtime     │ conda-forge      │ + pyvista, jupyter│
  │ dev         │ conda-forge      │ + claude, black   │
  ├─────────────┼──────────────────┼───────────────────┤
  │ amr         │ custom (~1hr)    │ AMR tools         │
  │ amr-runtime │ custom (~1hr)    │ + pyvista, jupyter│
  │ amr-dev     │ custom (~1hr)    │ + claude, black   │
  └─────────────┴──────────────────┴───────────────────┘

  MPI override environments (when you need the non-default MPI):
    mpich, mpich-dev, amr-mpich, amr-mpich-dev
    openmpi, openmpi-dev, amr-openmpi, amr-openmpi-dev

  Platform defaults: macOS → OpenMPI, Linux → MPICH.
  MPICH 4.x has a known scaling bug on macOS (issue #68).

COMMANDS

  Setup:
    setup               Interactive configuration wizard
    set-env NAME        Change environment directly
    ai-tools            Configure external AI instruction paths
    claude-perms        Configure Claude Code permissions (safe defaults)
    install-claude      Install Claude Code CLI into the dev pixi env (run via ./uw claude)

  Building:
    build               Build underworld3
    clean               Clean build artifacts
    petsc-local-build   Build custom PETSc + petsc4py (AMR, ~1 hour)
    petsc-local-clean   Remove custom PETSc build

  Testing:
    test [opts]         Level 1 tests (~2 min, fast mode)
    test-all [opts]     All test levels (~20 min)

    Options pass through to test script:
      --isolation       Run each file in subprocess (thorough)
      --parallel        Include MPI tests (2 ranks)
      --parallel-ranks N  MPI tests with N ranks
      1,2,3             Override test levels

    Examples:
      ./uw test                     Quick level 1, fast mode
      ./uw test --isolation         Level 1 with process isolation
      ./uw test --parallel 2        Level 2 with MPI tests
      ./uw test-all --isolation     All levels, isolated (CI mode)
      ./uw test --help              Show all test options

  Development:
    shell               Interactive pixi shell in the configured env (avoid bare 'pixi shell')
    jupyter [args]      Run jupyter (add --worktree NAME to open a worktree)

  Worktrees:
    worktree create <name> [prefix]   Create worktree (own .pixi, shared PETSc)
    worktree shell <name>             Start shell in worktree directory
    worktree list                     List worktrees with status
    worktree remove <name>            Remove worktree and branch

    Worktrees share the main repo's pixi environment and PETSc build.
    Once inside a worktree shell, ./uw build installs from that source.

  Diagnostics:
    doctor              Check for configuration issues (PETSc mismatch, etc.)
    status              Check for updates on GitHub (without pulling)
    update              Pull latest changes and rebuild
EOF
    if [ "$verbose" = "1" ]; then
        cat << 'EOF'

  Developer / maintenance  (./uw dev <tool>):
    dev release [opts]       Orchestrate development -> main promotion
                             (--version vX.Y.Z  --rc N  --yes  --dry-run)
    dev release-check [lvl]  Per-feature validation gate (dry-run; default 1,2)
    dev branch-hygiene       Report/prune branches merged to development
    dev feature-audit        List feature merges available to release
    dev format               Run black
    dev type-check           Run mypy

    See docs/developer/guides/release-process.md  (or ./uw dev --help)
EOF
    else
        cat << 'EOF'

  Developer tools live under ./uw dev — reveal them with:
    ./uw --help all          (shown by default on *-dev environments)
EOF
    fi
    cat << 'EOF'

  Any command not listed is passed to: pixi run -e $ENV <command>

FILES

  .pixi-env             Your environment choice (git-ignored)
  pixi.toml             Environment definitions (versioned)
  petsc-custom/petsc/   Custom PETSc build (git-ignored)

EOF
}

show_dev_help() {
    cat << 'EOF'
UNDERWORLD3 DEVELOPER TOOLS  —  ./uw dev <command>
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Maintainer and contributor tooling. These run in any environment, but the
runtime/default envs are deliberately stripped down for beginners, demos and
binder — so on those envs ./uw --help hides this section (reveal it with
./uw --help all). On a *-dev environment it is shown by default.

  Releasing (development -> main):
    dev release [opts]       Orchestrate the quarterly promotion: tag an RC on
                             development, run full validation + the per-feature
                             gate, scaffold release notes, open the dev->main PR.
                             Stops before the production tag (review gate intact).
                             Options: --version vX.Y.Z  --rc N  --yes  --dry-run
    dev release-check [lvl]  Per-feature validation gate, dry-run (default 1,2).
                             Tells you which features pass — i.e. are announceable
                             as "supported". Run on development.
    dev branch-hygiene       Report branches merged into development, branches
                             with open PRs, and stale branches. Add --prune to
                             delete merged branches (confirms each).
    dev feature-audit        List feature-branch PR merges into development since
                             the last release — the pool of features available to
                             announce. Map these to entries in the manifest.
                             [--all also lists bugfix/docs; --since <ref> overrides]

  Code quality:
    dev format               Run black over src/ and tests/
    dev type-check           Run mypy over src/

  Guide: docs/developer/guides/release-process.md

EOF
}

# Check for updates on remote (without pulling)
check_for_updates() {
    local quiet="${1:-false}"

    # Check if this is a git repo
    if [ ! -d "$SCRIPT_DIR/.git" ]; then
        [ "$quiet" = "false" ] && echo "Not a git repository"
        return 1
    fi

    cd "$SCRIPT_DIR"

    # Get current branch
    local branch=$(git rev-parse --abbrev-ref HEAD 2>/dev/null)
    if [ -z "$branch" ]; then
        [ "$quiet" = "false" ] && echo "Could not determine current branch"
        return 1
    fi

    # Fetch updates from remote (just metadata, no merge)
    if [ "$quiet" = "false" ]; then
        echo "Checking for updates on branch: $branch"
        echo ""
    fi

    git fetch origin "$branch" --quiet 2>/dev/null

    # Count commits behind/ahead
    local behind=$(git rev-list --count HEAD..origin/"$branch" 2>/dev/null || echo "0")
    local ahead=$(git rev-list --count origin/"$branch"..HEAD 2>/dev/null || echo "0")

    if [ "$quiet" = "true" ]; then
        # Return status for quiet mode
        [ "$behind" -gt 0 ] && return 0 || return 1
    fi

    # Show status
    if [ "$behind" -eq 0 ] && [ "$ahead" -eq 0 ]; then
        echo -e "  ${GREEN}✓${NC} Up to date with origin/$branch"
    else
        if [ "$behind" -gt 0 ]; then
            echo -e "  ${YELLOW}↓${NC} $behind commit(s) behind origin/$branch"
            echo ""
            echo "  Recent changes on remote:"
            git log --oneline HEAD..origin/"$branch" | head -5 | sed 's/^/    /'
            local total_behind=$(git rev-list --count HEAD..origin/"$branch")
            [ "$total_behind" -gt 5 ] && echo "    ... and $((total_behind - 5)) more"
        fi
        if [ "$ahead" -gt 0 ]; then
            echo -e "  ${GREEN}↑${NC} $ahead commit(s) ahead of origin/$branch"
        fi
        echo ""
        if [ "$behind" -gt 0 ]; then
            echo "  Run './uw update' to pull changes and rebuild"
        fi
    fi

    # Check for local modifications
    if ! git diff --quiet 2>/dev/null || ! git diff --cached --quiet 2>/dev/null; then
        echo ""
        echo -e "  ${YELLOW}!${NC} You have uncommitted local changes"
    fi

    return 0
}

# Pull updates and rebuild
run_update() {
    cd "$SCRIPT_DIR"

    # Check if this is a git repo
    if [ ! -d "$SCRIPT_DIR/.git" ]; then
        echo "Not a git repository"
        exit 1
    fi

    local branch=$(git rev-parse --abbrev-ref HEAD 2>/dev/null)
    echo -e "${BOLD}Updating Underworld3${NC} (branch: $branch)"
    echo ""

    # Check for local changes
    if ! git diff --quiet 2>/dev/null || ! git diff --cached --quiet 2>/dev/null; then
        echo -e "${YELLOW}You have uncommitted local changes:${NC}"
        git status --short
        echo ""
        read -p "Stash changes and continue? [y/N]: " stash_choice
        if [ "$stash_choice" = "y" ] || [ "$stash_choice" = "Y" ]; then
            git stash push -m "uw update: auto-stash $(date +%Y-%m-%d_%H:%M)"
            echo ""
        else
            echo "Commit or stash your changes first, then run './uw update'"
            exit 1
        fi
    fi

    # Pull updates
    echo "Pulling from origin/$branch..."
    if git pull origin "$branch"; then
        echo ""
        echo -e "${GREEN}✓${NC} Pull successful"
        echo ""

        # Rebuild
        read -p "Rebuild underworld3? [Y/n]: " rebuild_choice
        rebuild_choice=${rebuild_choice:-y}
        if [ "$rebuild_choice" = "y" ] || [ "$rebuild_choice" = "Y" ]; then
            echo ""
            run_build
        else
            echo ""
            echo "Run './uw build' when ready to rebuild"
        fi
    else
        echo ""
        echo -e "${YELLOW}Pull failed - check errors above${NC}"
        exit 1
    fi
}

# ── Worktree management ──────────────────────────────────────────────
#
# Each worktree gets its OWN pixi environment (isolated site-packages).
# Only PETSc is shared via symlink (non-relocatable, expensive to rebuild).
# `./uw build` from inside a worktree installs into that worktree's env.

WORKTREE_ROOT="$SCRIPT_DIR/.claude/worktrees"

# Resolve the main repo root.  If we ARE a worktree, follow the chain
# back to the principal working tree.
get_main_repo() {
    git -C "$SCRIPT_DIR" worktree list --porcelain 2>/dev/null \
        | head -1 | sed 's/^worktree //'
}

worktree_create() {
    local name="$1"
    local branch_prefix="${2:-feature}"   # feature, bugfix, docs, …

    if [ -z "$name" ]; then
        echo "Usage: ./uw worktree create <name> [branch-prefix]"
        echo ""
        echo "  branch-prefix: feature (default), bugfix, docs, etc."
        echo ""
        echo "Examples:"
        echo "  ./uw worktree create viscoelasticity"
        echo "  ./uw worktree create mesh-fix bugfix"
        exit 1
    fi

    # Policy: worktrees must NEVER be on `development` or `main` directly.
    # `main` is the release branch (tagged history); `development` is the
    # integration trunk and lives in the default repo checkout. All work
    # happens on side branches. Refuse misleading prefixes.
    if [ "$branch_prefix" = "development" ] || [ "$branch_prefix" = "main" ]; then
        local main_repo
        main_repo=$(get_main_repo)
        echo -e "${YELLOW}Refusing to use '$branch_prefix' as a branch prefix.${NC}"
        echo ""
        echo "Worktrees in this repo are always on a side branch:"
        echo "  feature/<name>   bugfix/<name>   docs/<name>   ..."
        echo ""
        echo "If you wanted a checkout of '$branch_prefix' to read code from,"
        echo "use the main repo at:"
        echo "  $main_repo"
        echo ""
        echo "If you wanted to do work that will merge to development,"
        echo "use a regular prefix:"
        echo "  ./uw worktree create $name           # → feature/$name"
        echo "  ./uw worktree create $name bugfix    # → bugfix/$name"
        exit 1
    fi

    local main_repo=$(get_main_repo)
    local wt_path="$main_repo/.claude/worktrees/$name"
    local branch_name="${branch_prefix}/${name}"

    if [ -d "$wt_path" ]; then
        echo "Worktree '$name' already exists at $wt_path"
        exit 1
    fi

    echo -e "${BOLD}Creating worktree:${NC} $name"
    echo "  Path:   $wt_path"
    echo "  Branch: $branch_name"
    echo ""

    # 1. Create the worktree
    mkdir -p "$main_repo/.claude/worktrees"
    git -C "$main_repo" worktree add "$wt_path" -b "worktree-$name" 2>/dev/null || {
        # Branch may already exist
        git -C "$main_repo" worktree add "$wt_path" "worktree-$name" 2>/dev/null || {
            echo -e "${YELLOW}Failed to create worktree${NC}"
            exit 1
        }
    }

    # 2. Reset to development (worktrees branch from HEAD which may be main)
    echo "  Resetting to origin/development..."
    git -C "$wt_path" fetch origin
    git -C "$wt_path" reset --hard origin/development

    # 3. Rename branch to follow convention
    git -C "$wt_path" branch -m "worktree-$name" "$branch_name" 2>/dev/null || true

    # 4. Set up isolated environment
    #
    # Each worktree gets its OWN pixi environment (not a symlink to the main
    # repo's .pixi). This prevents Cython binary incompatibilities when
    # worktrees have different source code. Only PETSc is shared (it's
    # non-relocatable and expensive to rebuild).
    echo "  Setting up isolated environment..."

    # petsc-custom/petsc → main repo's PETSc build (read-only, shared)
    if [ -d "$main_repo/petsc-custom/petsc" ]; then
        ln -sf "$main_repo/petsc-custom/petsc" "$wt_path/petsc-custom/petsc"
        echo -e "    ${GREEN}✓${NC} petsc-custom/petsc → main repo (shared)"
    fi

    # .petsc-version — copy (not symlink) so each worktree can select a
    # different already-built PETSC_ARCH (e.g. petsc-324-uw-openmpi vs
    # petsc-325-uw-openmpi).  Built arch directories coexist in the
    # shared PETSc source tree; .petsc-version just picks which one
    # the activation script uses for PETSC_ARCH.
    if [ -f "$main_repo/petsc-custom/.petsc-version" ]; then
        cp "$main_repo/petsc-custom/.petsc-version" "$wt_path/petsc-custom/.petsc-version"
        echo -e "    ${GREEN}✓${NC} .petsc-version copied ($(cat "$main_repo/petsc-custom/.petsc-version"))"
    fi

    # .pixi-env — copy so ./uw knows which environment to use
    if [ -f "$main_repo/.pixi-env" ]; then
        cp "$main_repo/.pixi-env" "$wt_path/.pixi-env"
        echo -e "    ${GREEN}✓${NC} .pixi-env copied ($(cat "$main_repo/.pixi-env"))"
    fi

    # Install pixi environment (own copy — NEVER symlink)
    local env=$(cat "$wt_path/.pixi-env" 2>/dev/null || echo "runtime")
    echo "  Installing pixi environment ($env) — this may take a moment..."
    (cd "$wt_path" && $PIXI install -e "$env" 2>/dev/null) && \
        echo -e "    ${GREEN}✓${NC} .pixi/ installed (isolated)" || {
        echo -e "    ${RED}pixi install failed for worktree${NC}"
        echo -e "    The worktree requires its own pixi environment."
        echo -e "    Fix the pixi install error above, then run:"
        echo -e "      cd $wt_path && pixi install -e $env"
        exit 1
    }

    # 4b. Install petsc4py for AMR environments
    # petsc4py is built from source (not a conda package), so pixi install
    # doesn't include it. We install from the shared PETSc build.
    # IMPORTANT: Don't cd into the petsc source dir — pixi resolves the
    # project from the real working directory, so cd'ing through a symlink
    # into the main repo makes pixi install into the wrong environment.
    if is_amr_env "$env" && [ -d "$wt_path/petsc-custom/petsc/src/binding/petsc4py" ]; then
        echo "  Installing petsc4py for $env..."
        rm -rf "$wt_path/petsc-custom/petsc/src/binding/petsc4py/build"   # shared stale-ABI cache
        (cd "$wt_path" && $PIXI run -e "$env" pip install \
            "$wt_path/petsc-custom/petsc/src/binding/petsc4py" --no-build-isolation 2>/dev/null) && \
            echo -e "    ${GREEN}✓${NC} petsc4py installed" || \
            echo -e "    ${YELLOW}⚠${NC} petsc4py install failed — run ./uw build from the worktree"
    fi

    # 5. Exclude local files from git tracking in this worktree.
    local git_dir
    git_dir=$(git -C "$wt_path" rev-parse --git-dir)
    local exclude_file="$git_dir/info/exclude"
    mkdir -p "$(dirname "$exclude_file")"
    for entry in .pixi petsc-custom/petsc .pixi-env; do
        if ! grep -qxF "$entry" "$exclude_file" 2>/dev/null; then
            echo "$entry" >> "$exclude_file"
        fi
    done
    echo -e "    ${GREEN}✓${NC} git exclude updated"

    echo ""
    echo -e "${GREEN}${BOLD}Worktree ready!${NC}"
    echo ""
    echo "  Start working:  ./uw worktree shell $name"
    echo "  Build:           cd $wt_path && ./uw build"
    echo "  Remove:          ./uw worktree remove $name"
}

worktree_shell() {
    local name="$1"

    if [ -z "$name" ]; then
        echo "Usage: ./uw worktree shell <name>"
        echo ""
        worktree_list
        exit 1
    fi

    local main_repo=$(get_main_repo)
    local wt_path="$main_repo/.claude/worktrees/$name"

    if [ ! -d "$wt_path" ]; then
        echo "Worktree '$name' not found."
        echo ""
        worktree_list
        exit 1
    fi

    local branch=$(git -C "$wt_path" rev-parse --abbrev-ref HEAD 2>/dev/null || echo "unknown")
    local env=$(cat "$wt_path/.pixi-env" 2>/dev/null || get_env)

    echo -e "${BOLD}Entering worktree:${NC} $name"
    echo "  Branch: $branch"
    echo "  Env:    $env"
    echo "  Path:   $wt_path"
    echo ""
    echo "  ./uw build   — build from this worktree's source"
    echo "  ./uw test    — run tests"
    echo "  ./uw shell   — pixi shell (if you need raw access)"
    echo "  exit         — leave worktree shell"
    echo ""

    # Start a pixi shell cd'd into the worktree
    cd "$wt_path"
    exec $PIXI shell -e "$env"
}

worktree_list() {
    local main_repo=$(get_main_repo)
    local wt_root="$main_repo/.claude/worktrees"

    echo -e "${BOLD}Worktrees:${NC}"
    echo ""

    if [ ! -d "$wt_root" ]; then
        echo "  (none)"
        echo ""
        echo "Create one:  ./uw worktree create <name>"
        return
    fi

    local found=0
    for d in "$wt_root"/*/; do
        [ ! -d "$d" ] && continue
        local wt_name=$(basename "$d")

        if [ -f "$d/.git" ]; then
            local branch=$(git -C "$d" rev-parse --abbrev-ref HEAD 2>/dev/null || echo "?")
            local status=$(git -C "$d" status --short 2>/dev/null | wc -l | tr -d ' ')
            local linked=$( [ -L "$d/.pixi" ] && echo "shared env" || echo "isolated" )
            local dirty=""
            [ "$status" -gt 0 ] && dirty="  ${YELLOW}($status modified)${NC}"

            # Policy advisory: worktrees should never be on development or main
            local policy=""
            if [ "$branch" = "development" ] || [ "$branch" = "main" ]; then
                policy="  ${YELLOW}⚠ on '$branch' — branch off to a side branch${NC}"
            fi

            echo -e "  ${GREEN}${wt_name}${NC}  [$branch]  ($linked)$dirty$policy"
            echo "    $d"
        else
            echo -e "  ${YELLOW}${wt_name}${NC}  (not a git worktree)"
        fi
        found=1
    done

    [ "$found" -eq 0 ] && echo "  (none)"
    echo ""
    echo "Commands:"
    echo "  ./uw worktree create <name>   Create a new worktree"
    echo "  ./uw worktree shell <name>    Start a shell in the worktree"
    echo "  ./uw worktree remove <name>   Remove a worktree"
}

worktree_remove() {
    local name="$1"

    if [ -z "$name" ]; then
        echo "Usage: ./uw worktree remove <name>"
        echo ""
        worktree_list
        exit 1
    fi

    local main_repo=$(get_main_repo)
    local wt_path="$main_repo/.claude/worktrees/$name"

    if [ ! -d "$wt_path" ]; then
        echo "Worktree '$name' not found."
        exit 1
    fi

    local branch=$(git -C "$wt_path" rev-parse --abbrev-ref HEAD 2>/dev/null || echo "")

    # Filter out symlinks we created (.pixi, petsc-custom/petsc) from the dirty check
    local status=$(git -C "$wt_path" status --short 2>/dev/null \
        | grep -v '^?? \.pixi$' \
        | grep -v '^?? petsc-custom/petsc$' \
        | wc -l | tr -d ' ')

    if [ "$status" -gt 0 ]; then
        echo -e "${YELLOW}Worktree has $status uncommitted change(s) (excluding setup symlinks):${NC}"
        git -C "$wt_path" status --short \
            | grep -v '^?? \.pixi$' \
            | grep -v '^?? petsc-custom/petsc$'
        echo ""
        read -p "Remove anyway? [y/N]: " confirm
        [ "$confirm" != "y" ] && [ "$confirm" != "Y" ] && exit 1
    fi

    echo "Removing worktree '$name'..."

    # Must not be cd'd into the worktree
    cd "$main_repo"

    git worktree remove "$wt_path" --force 2>/dev/null || {
        echo -e "${YELLOW}git worktree remove failed — removing manually${NC}"
        rm -rf "$wt_path"
        git worktree prune
    }

    # Delete the branch (use -D because it's merged to development, not main)
    if [ -n "$branch" ] && [ "$branch" != "HEAD" ] && [ "$branch" != "development" ]; then
        git branch -D "$branch" 2>/dev/null && \
            echo -e "  ${GREEN}✓${NC} Deleted branch $branch" || true
    fi

    echo -e "${GREEN}✓${NC} Worktree '$name' removed"
}

# ── Release tooling ─────────────────────────────────────────────────────────
# Maturity-gated quarterly promotion of development -> main.
# See docs/developer/guides/release-process.md

# Small y/N confirm helper. Returns 0 on yes. Honours an assume-yes flag.
_confirm() {
    local prompt="$1" assume_yes="${2:-0}"
    [ "$assume_yes" -eq 1 ] && return 0
    read -p "$prompt [y/N]: " _c
    [ "$_c" = "y" ] || [ "$_c" = "Y" ]
}

# ./uw dev release-check [levels]
# Dry-run of the per-feature validation gate against the current checkout.
# Read-only; tells you which features currently pass (announceable as supported).
run_release_check() {
    local levels="${1:-1,2}"
    cd "$SCRIPT_DIR"
    local branch=$(git rev-parse --abbrev-ref HEAD 2>/dev/null)
    echo -e "${BOLD}Release validation gate${NC}  (branch: $branch, levels: $levels)"
    if [ "$branch" != "development" ]; then
        echo -e "  ${YELLOW}note:${NC} normally run on 'development'; you are on '$branch'."
    fi
    $PIXI run -e "$(get_env)" python "$SCRIPT_DIR/scripts/release_gate.py" \
        --levels "$levels" report
}

# ./uw dev branch-hygiene [--prune]
# Report (and optionally delete) branches that are merged into development.
# Never touches branches with an open PR. Targets the "dangling branches" pileup.
run_branch_hygiene() {
    local do_prune=0
    [ "$1" = "--prune" ] && do_prune=1
    cd "$SCRIPT_DIR"

    if [ ! -d "$SCRIPT_DIR/.git" ]; then
        echo "Not a git repository"; exit 1
    fi

    echo -e "${BOLD}Fetching and pruning remote-tracking refs...${NC}"
    git fetch --prune --quiet origin 2>/dev/null || git fetch --prune origin

    # Head branches of open PRs — never prune these.
    # gh_ok is set only when the open-PR query genuinely succeeds. An empty list
    # from a *successful* query means "no open PRs" (safe). An empty list from a
    # failed/unauth'd query must NOT be treated as "no open PRs" — otherwise
    # --prune could delete branches that actually have open PRs.
    local open_pr_branches="" gh_ok=0
    if command -v gh >/dev/null 2>&1; then
        if open_pr_branches=$(gh pr list --state open --limit 200 \
            --json headRefName --jq '.[].headRefName' 2>/dev/null); then
            gh_ok=1
        else
            echo -e "  ${YELLOW}gh could not list open PRs (auth/access?) — pruning disabled.${NC}"
        fi
    else
        echo -e "  ${YELLOW}gh not found — cannot cross-check open PRs; pruning disabled.${NC}"
    fi

    # Remote branches merged into development (prune candidates), minus the trunks.
    local merged
    merged=$(git branch -r --merged origin/development 2>/dev/null \
        | sed 's/^[* ]*//' \
        | grep '^origin/' \
        | grep -vE 'origin/(HEAD|development|main)$' \
        | sed 's#^origin/##' \
        | sort -u)

    echo ""
    echo -e "${BOLD}Merged into development (safe to delete):${NC}"
    local prunable=()
    if [ -z "$merged" ]; then
        echo "  (none)"
    else
        while IFS= read -r b; do
            [ -z "$b" ] && continue
            if echo "$open_pr_branches" | grep -qx "$b"; then
                echo -e "  ${YELLOW}$b${NC}  (has open PR — keeping)"
            else
                echo -e "  ${GREEN}$b${NC}"
                prunable+=("$b")
            fi
        done <<< "$merged"
    fi

    echo ""
    echo -e "${BOLD}Open PRs (do NOT prune):${NC}"
    if [ -z "$open_pr_branches" ]; then
        echo "  (none)"
    else
        echo "$open_pr_branches" | sed 's/^/  /'
    fi

    # Stale: not merged, no open PR, last commit older than 90 days.
    echo ""
    echo -e "${BOLD}Stale branches (unmerged, no open PR, >90 days):${NC}"
    local now=$(date +%s)
    local cutoff=$(( 90 * 24 * 3600 ))
    local found_stale=0
    while IFS='|' read -r b ts; do
        [ -z "$b" ] && continue
        b=${b#origin/}
        case "$b" in HEAD|development|main) continue ;; esac
        echo "$merged" | grep -qx "$b" && continue
        echo "$open_pr_branches" | grep -qx "$b" && continue
        local age=$(( now - ts ))
        if [ "$age" -gt "$cutoff" ]; then
            echo -e "  ${YELLOW}$b${NC}  ($(( age / 86400 ))d)"
            found_stale=1
        fi
    done <<< "$(git for-each-ref --format='%(refname:short)|%(committerdate:unix)' refs/remotes/origin)"
    [ "$found_stale" -eq 0 ] && echo "  (none)"

    echo ""
    if [ "$do_prune" -eq 0 ]; then
        echo "Report only. Re-run with --prune to delete merged branches:"
        echo "  ./uw dev branch-hygiene --prune"
        return
    fi

    if [ "$gh_ok" -eq 0 ]; then
        echo -e "${YELLOW}Refusing to --prune: the open-PR cross-check is unavailable,${NC}"
        echo "so a branch with an open PR could be deleted by mistake."
        echo "Authenticate gh ('gh auth login') and retry."
        return
    fi

    if [ "${#prunable[@]}" -eq 0 ]; then
        echo "Nothing to prune."
        return
    fi

    echo -e "${BOLD}Prune merged branches:${NC}"
    local b
    for b in "${prunable[@]}"; do
        read -p "  Delete origin/$b (and local if present)? [y/N]: " confirm
        if [ "$confirm" = "y" ] || [ "$confirm" = "Y" ]; then
            git push origin --delete "$b" 2>/dev/null \
                && echo -e "    ${GREEN}✓${NC} deleted origin/$b" \
                || echo -e "    ${YELLOW}!${NC} could not delete origin/$b"
            git branch -D "$b" 2>/dev/null \
                && echo -e "    ${GREEN}✓${NC} deleted local $b" || true
        fi
    done
}

# ./uw dev feature-audit [--all] [--since <ref>]
# List feature-branch PR merges into development since the last release, so you
# know what is available to announce. Candidates for the release manifest.
run_feature_audit() {
    local show_all=0 since=""
    while [ $# -gt 0 ]; do
        case "$1" in
            --all)   show_all=1; shift ;;
            --since) since="$2"; shift 2 ;;
            *) echo "Usage: ./uw dev feature-audit [--all] [--since <ref>]"; exit 1 ;;
        esac
    done
    cd "$SCRIPT_DIR"
    git fetch --tags --quiet origin 2>/dev/null

    # Baseline: explicit --since, else the last release tag if it is on
    # development, else the point where development last diverged from main.
    local base base_label
    if [ -n "$since" ]; then
        base="$since"; base_label="$since"
    else
        local last_tag=$(git tag --list 'v*' --sort=-v:refname | head -1)
        if [ -n "$last_tag" ] && git merge-base --is-ancestor "$last_tag" origin/development 2>/dev/null; then
            base="$last_tag"; base_label="release $last_tag"
        else
            base=$(git merge-base origin/main origin/development)
            base_label="$(git rev-parse --short "$base") (main divergence point)"
        fi
    fi
    local base_date=$(git log -1 --format=%cd --date=short "$base" 2>/dev/null)

    echo -e "${BOLD}Feature merges into development since ${base_label}, ${base_date}:${NC}"
    echo "  candidates for docs/release-notes/feature-manifest.yaml"
    echo ""

    local pattern='from [^ ]*/feature/'
    [ "$show_all" -eq 1 ] && pattern='from [^ ]*/(feature|bugfix|docs)/'

    local count=0
    while IFS='|' read -r d s; do
        [ -z "$s" ] && continue
        local pr=$(echo "$s" | grep -oE '#[0-9]+' | head -1)
        local br=$(echo "$s" | sed -E 's#.*from [^/ ]+/##')
        printf "  ${GREEN}%-12s${NC} %-7s %s\n" "$d" "$pr" "$br"
        count=$((count + 1))
    done <<< "$(git log "$base"..origin/development --merges --first-parent \
        --date=short --format='%cd|%s' 2>/dev/null | grep -E "$pattern")"

    local listed=$(grep -c '^  - key:' docs/release-notes/feature-manifest.yaml 2>/dev/null || echo 0)
    echo ""
    echo "  $count merge(s) found. Manifest currently lists $listed feature(s)."
    echo "  Note: only PR-merge commits are detected; squash-merged PRs won't appear."
    echo "  Validated-test coverage:  python scripts/check_manifest_coverage.py"
}

# ./uw dev release [--version vX.Y.Z] [--rc N] [--yes] [--dry-run]
# Orchestrate the quarterly development -> main promotion. Every state-changing
# step is confirmed. Stops BEFORE the production tag push so main's branch
# protection (PR + review) stays intact; CI release.yml owns the GitHub release.
run_release() {
    local version="" rc="1" assume_yes=0 dry_run=0
    while [ $# -gt 0 ]; do
        case "$1" in
            --version) version="$2"; shift 2 ;;
            --rc) rc="$2"; shift 2 ;;
            --yes|-y) assume_yes=1; shift ;;
            --dry-run) dry_run=1; shift ;;
            *) echo "Unknown option: $1"; echo "Usage: ./uw dev release [--version vX.Y.Z] [--rc N] [--yes] [--dry-run]"; exit 1 ;;
        esac
    done

    cd "$SCRIPT_DIR"

    echo -e "${BOLD}[1/7] Preflight${NC}"
    local branch=$(git rev-parse --abbrev-ref HEAD 2>/dev/null)
    if [ "$branch" != "development" ]; then
        echo -e "  ${YELLOW}Release must be cut from 'development' in the main checkout (you are on '$branch').${NC}"
        exit 1
    fi
    if [ -n "$(git status --porcelain)" ]; then
        echo -e "  ${YELLOW}Working tree is not clean. Commit or stash first.${NC}"
        exit 1
    fi
    git fetch --quiet origin 2>/dev/null || git fetch origin
    local behind=$(git rev-list --count HEAD..origin/development 2>/dev/null || echo 0)
    if [ "$behind" -gt 0 ]; then
        echo -e "  ${YELLOW}Local development is $behind commit(s) behind origin. Pull first.${NC}"; exit 1
    fi
    local unpushed=$(git rev-list --count origin/development..HEAD 2>/dev/null || echo 0)
    if [ "$unpushed" -gt 0 ]; then
        echo -e "  ${YELLOW}Local development has $unpushed unpushed commit(s). Push first so the RC tag matches what reviewers and CI see on origin.${NC}"; exit 1
    fi
    local ahead=$(git rev-list --count origin/main..origin/development 2>/dev/null || echo 0)
    echo -e "  ${GREEN}✓${NC} on development, clean, in sync with origin, $ahead commit(s) ahead of main"

    echo ""
    echo -e "${BOLD}[2/7] Version${NC}"
    if [ -z "$version" ]; then
        local last_tag=$(git tag --list 'v*' --sort=-v:refname | head -1)
        echo "  Last release tag: ${last_tag:-none}"
        read -p "  Release version (e.g. v3.1.0): " version
    fi
    [ -z "$version" ] && { echo "No version given."; exit 1; }
    case "$version" in v*) ;; *) version="v$version" ;; esac
    local rc_tag="${version}rc${rc}"
    echo "  Release:  $version"
    echo "  RC tag:   $rc_tag  (pushed to development now; production tag is manual after merge)"

    echo ""
    echo -e "${BOLD}[3/7] Tag release candidate on development${NC}"
    if [ "$dry_run" -eq 1 ]; then
        echo "  (dry-run) would: git tag $rc_tag && git push origin $rc_tag"
    elif _confirm "  Tag and push $rc_tag on development?" "$assume_yes"; then
        git tag "$rc_tag" && git push origin "$rc_tag" \
            && echo -e "  ${GREEN}✓${NC} pushed $rc_tag (CI builds a prerelease)"
    else
        echo "  Skipped RC tag."
    fi

    echo ""
    echo -e "${BOLD}[4/7] Full validation + per-feature gate${NC}"
    local results_json="${TMPDIR:-/tmp}/uw_release_gate_${version}.json"
    if [ "$dry_run" -eq 1 ]; then
        echo "  (dry-run) would run: ./scripts/test_levels.sh 1,2,3 --isolation"
        echo "  (dry-run) would run: release_gate.py --levels 1,2,3 run > $results_json"
    else
        echo "  Running full test suite (levels 1,2,3, isolated)..."
        $PIXI run -e "$(get_env)" ./scripts/test_levels.sh 1,2,3 --isolation || \
            echo -e "  ${YELLOW}! Test suite reported failures — review before publishing.${NC}"
        echo "  Running per-feature maturity gate..."
        $PIXI run -e "$(get_env)" python "$SCRIPT_DIR/scripts/release_gate.py" \
            --levels 1,2,3 run > "$results_json"
        $PIXI run -e "$(get_env)" python "$SCRIPT_DIR/scripts/release_gate.py" \
            --results "$results_json" report
    fi

    echo ""
    echo -e "${BOLD}[5/7] Release notes${NC}"
    local notes_file="docs/release-notes/${version}.md"
    if [ "$dry_run" -eq 1 ]; then
        echo "  (dry-run) would render notes to $notes_file"
    elif [ -f "$notes_file" ]; then
        echo -e "  ${YELLOW}$notes_file already exists — not overwriting.${NC}"
    elif _confirm "  Generate $notes_file from the gate results?" "$assume_yes"; then
        $PIXI run -e "$(get_env)" python "$SCRIPT_DIR/scripts/release_gate.py" \
            --results "$results_json" render-notes "$version" > "$notes_file"
        echo -e "  ${GREEN}✓${NC} wrote $notes_file — edit Highlights & Contributors, then commit it."
    fi

    echo ""
    echo -e "${BOLD}[6/7] Open development -> main PR${NC}"
    if [ "$dry_run" -eq 1 ]; then
        echo "  (dry-run) would: gh pr create --base main --head development ..."
    elif ! command -v gh >/dev/null 2>&1; then
        echo -e "  ${YELLOW}gh not found — open the development->main PR manually.${NC}"
    elif _confirm "  Create the development -> main PR now?" "$assume_yes"; then
        if [ -f "$notes_file" ]; then
            gh pr create --base main --head development \
                --title "Release $version" --body-file "$notes_file" \
                && echo -e "  ${GREEN}✓${NC} PR opened"
        else
            gh pr create --base main --head development \
                --title "Release $version" --body "Release $version" \
                && echo -e "  ${GREEN}✓${NC} PR opened"
        fi
    fi

    echo ""
    echo -e "${BOLD}[7/7] Final tag (manual, after the PR merges to main)${NC}"
    echo "  Once the development -> main PR is reviewed and merged:"
    echo "    git checkout main && git pull"
    echo "    git tag $version && git push origin $version"
    echo "  CI (release.yml) then publishes the GitHub Release from $notes_file."
    echo ""
    echo -e "${GREEN}Done.${NC} The production tag is deliberately left for you to push"
    echo "after the merge, so main's PR + review gate stays intact."
}

# Check if an environment is installed
env_installed() {
    [ -d "$SCRIPT_DIR/.pixi/envs/$1" ]
}

# Check if underworld3 is built in an environment
uw3_built() {
    [ -d "$SCRIPT_DIR/.pixi/envs/$1/lib/python3.12/site-packages/underworld3" ]
}

# Show environment status table
show_env_status() {
    local current_env=$(get_env)
    local has_petsc=$(petsc_built && echo "yes" || echo "no")

    echo "Environments:  (* = selected,  [✓] = uw3 built,  [ ] = needs ./uw build)"
    echo ""

    # Determine platform-default MPI for display
    local default_mpi="mpich"
    [[ "$(uname -s)" == "Darwin" ]] && default_mpi="openmpi"

    for env in default runtime dev amr amr-runtime amr-dev \
               mpich mpich-dev amr-mpich amr-mpich-dev \
               openmpi openmpi-dev amr-openmpi amr-openmpi-dev; do
        # Only show environments that are installed or selected
        if ! env_installed "$env" && [ "$env" != "$current_env" ]; then
            continue
        fi

        # Selected marker
        local marker="  "
        [ "$env" = "$current_env" ] && marker="* "

        # Build status
        local build_status="[ ]"
        if ! env_installed "$env"; then
            build_status="───"  # Not installed
        elif uw3_built "$env"; then
            build_status="[✓]"
        fi

        # Description
        local desc=""
        case "$env" in
            default)         desc="minimal ($default_mpi)" ;;
            runtime)         desc="+ viz, jupyter ($default_mpi)" ;;
            dev)             desc="+ claude, linting ($default_mpi)" ;;
            mpich)           desc="minimal (mpich override)" ;;
            mpich-dev)       desc="+ claude, linting (mpich override)" ;;
            openmpi)         desc="minimal (openmpi override)" ;;
            openmpi-dev)     desc="+ claude, linting (openmpi override)" ;;
            amr)             desc="custom PETSc ($default_mpi)" ;;
            amr-runtime)     desc="custom PETSc + viz ($default_mpi)" ;;
            amr-dev)         desc="custom PETSc + dev ($default_mpi)" ;;
            amr-mpich)       desc="custom PETSc (mpich override)" ;;
            amr-mpich-dev)   desc="custom PETSc + dev (mpich override)" ;;
            amr-openmpi)     desc="custom PETSc (openmpi override)" ;;
            amr-openmpi-dev) desc="custom PETSc + dev (openmpi override)" ;;
        esac

        # AMR note if PETSc not built for this env's arch
        if is_amr_env "$env" && ! petsc_built "$env"; then
            desc="$desc (petsc not built)"
        fi

        printf "  %s%-22s  %s  %s\n" "$marker" "$env" "$build_status" "$desc"
    done
}

# Brief welcome + offer setup
show_welcome() {
    echo -e "${BOLD}Underworld3${NC}"
    echo ""

    # Check for pixi first
    if [ -z "$PIXI" ]; then
        echo "Environment: (pixi not installed)"
        echo ""
        read -p "Run setup to install pixi? [Y/n]: " do_setup
        do_setup=${do_setup:-y}
        if [ "$do_setup" = "y" ] || [ "$do_setup" = "Y" ]; then
            echo ""
            run_setup
        else
            echo ""
            echo "Install pixi from https://pixi.sh then run './uw setup'"
        fi
        return
    fi

    if [ -f "$CONFIG_FILE" ]; then
        show_env_status
        echo ""
        echo "Commands:  ./uw setup   ./uw build   ./uw test   ./uw doctor   ./uw --help"
    else
        echo "Environment: (not configured)"
        echo ""
        read -p "Run setup? [Y/n]: " do_setup
        do_setup=${do_setup:-y}
        if [ "$do_setup" = "y" ] || [ "$do_setup" = "Y" ]; then
            echo ""
            run_setup
        else
            echo ""
            echo "Run './uw setup' when ready, or './uw --help' for options."
        fi
    fi
}

# Main
case "${1:-}" in
    "")
        show_welcome
        ;;
    setup)
        run_setup
        ;;
    set-env)
        if [ -z "$2" ]; then
            echo "Usage: ./uw set-env <environment>"
            echo ""
            echo "Conda-forge PETSc:"
            echo "  MPICH:   default, runtime, dev"
            echo "  OpenMPI: openmpi, openmpi-runtime, openmpi-dev"
            echo ""
            echo "Custom PETSc (AMR):"
            echo "  MPICH:   amr, amr-runtime, amr-dev"
            echo "  OpenMPI: amr-openmpi, amr-openmpi-runtime, amr-openmpi-dev"
            echo ""
            echo "Or run './uw setup' for interactive configuration."
            exit 1
        fi
        echo "$2" > "$CONFIG_FILE"
        echo "Environment: $2"
        echo "Run './uw build' to rebuild underworld3"
        ;;
    shell)
        # Activate the env stored in .pixi-env, not pixi's literal 'default'.
        # Bare `pixi shell` ignores .pixi-env and falls back to the env named
        # 'default' in pixi.toml, which is the lightweight conda-forge PETSc
        # env without nodejs / dev tools — a common first-timer trap.
        shift
        exec $PIXI shell -e "$(get_env)" "$@"
        ;;
    jupyter)
        shift
        # Check for --worktree flag
        WORKTREE_DIR=""
        JUPYTER_ARGS=()
        while [ $# -gt 0 ]; do
            case "$1" in
                --worktree)
                    if [ -z "$2" ]; then
                        echo "Usage: ./uw jupyter lab --worktree <name>"
                        echo "Run './uw worktrees' to see available worktrees."
                        exit 1
                    fi
                    WT_NAME="$2"
                    WT_PATH="$SCRIPT_DIR/.claude/worktrees/$WT_NAME"
                    if [ ! -d "$WT_PATH" ]; then
                        echo "Worktree '$WT_NAME' not found at $WT_PATH"
                        echo ""
                        echo "Available worktrees:"
                        if [ -d "$SCRIPT_DIR/.claude/worktrees" ]; then
                            for d in "$SCRIPT_DIR/.claude/worktrees"/*/; do
                                [ -d "$d" ] && echo "  $(basename "$d")"
                            done
                        else
                            echo "  (none)"
                        fi
                        exit 1
                    fi
                    WORKTREE_DIR="$WT_PATH"
                    shift 2
                    ;;
                *)
                    JUPYTER_ARGS+=("$1")
                    shift
                    ;;
            esac
        done
        if [ -n "$WORKTREE_DIR" ]; then
            echo -e "Starting JupyterLab in worktree: ${GREEN}$(basename "$WORKTREE_DIR")${NC}"
            echo "  Path: $WORKTREE_DIR"
            echo ""
            exec $PIXI run -e "$(get_env)" jupyter "${JUPYTER_ARGS[@]}" --notebook-dir="$WORKTREE_DIR"
        else
            exec $PIXI run -e "$(get_env)" jupyter "${JUPYTER_ARGS[@]}"
        fi
        ;;
    worktree)
        shift
        case "${1:-list}" in
            create)  shift; worktree_create "$@" ;;
            shell)   shift; worktree_shell "$@" ;;
            list)    worktree_list ;;
            remove)  shift; worktree_remove "$@" ;;
            *)
                echo "Usage: ./uw worktree <command>"
                echo ""
                echo "Commands:"
                echo "  create <name> [prefix]   Create a new worktree"
                echo "  shell <name>             Start a shell in the worktree"
                echo "  list                     List worktrees (default)"
                echo "  remove <name>            Remove a worktree"
                ;;
        esac
        ;;
    worktrees)
        # Backward compat alias
        worktree_list
        ;;
    build)
        run_build
        ;;
    clean)
        echo "Cleaning build artifacts..."
        rm -rf "$SCRIPT_DIR"/build/lib.* "$SCRIPT_DIR"/build/temp.* "$SCRIPT_DIR"/build/bdist.*
        rm -f "$SCRIPT_DIR/build/.petsc_target"
        echo -e "${GREEN}Done${NC}. Run './uw build' to rebuild."
        ;;
    petsc)
        # PETSc version management: ./uw petsc [versions|switch|active]
        run_petsc_cmd "$2" "$3"
        ;;
    doctor)
        run_doctor
        ;;
    status)
        check_for_updates
        ;;
    update)
        run_update
        ;;
    ai-tools)
        configure_ai_tools
        ;;
    claude-perms)
        configure_claude_permissions "$(get_env)"
        ;;
    install-claude)
        # The install-claude pixi task lives under [feature.dev.tasks], so it
        # only resolves in dev-feature envs (dev, amr-dev, *-dev). Guard the
        # call so users on default/runtime get a clear message instead of a
        # raw "task not found" error from pixi.
        env=$(get_env)
        case "$env" in
            *dev*)
                $PIXI run -e "$env" install-claude
                ;;
            *)
                echo -e "${YELLOW}install-claude needs a dev environment (current: $env).${NC}"
                echo "Switch with one of:"
                echo "  ./uw set-env dev          # conda-forge PETSc"
                echo "  ./uw set-env amr-dev      # custom PETSc with AMR tools"
                echo "  ./uw set-env mpich-dev    # MPICH override"
                echo "  ./uw set-env openmpi-dev  # OpenMPI override"
                exit 1
                ;;
        esac
        ;;
    dev)
        shift
        cmd="${1:-help}"
        case "$cmd" in
            help|--help|-h|"")
                show_dev_help
                ;;
            *)
                # Soft note, not a block: dev tools run in any environment, but
                # runtime/default are the stripped-down beginner/binder/teaching
                # envs, so flag when a developer tool is used there.
                case "$(get_env)" in
                    *dev*) ;;
                    *) echo -e "${YELLOW}note:${NC} 'dev $cmd' is a developer tool; you're on '$(get_env)' (a user env). Running anyway." >&2 ;;
                esac
                shift
                case "$cmd" in
                    release)        run_release "$@" ;;
                    release-check)  run_release_check "$@" ;;
                    branch-hygiene) run_branch_hygiene "$@" ;;
                    feature-audit)  run_feature_audit "$@" ;;
                    format)         exec $PIXI run -e "$(get_env)" format "$@" ;;
                    type-check)     exec $PIXI run -e "$(get_env)" type-check "$@" ;;
                    *)
                        echo "Unknown dev command: $cmd"
                        echo ""
                        show_dev_help
                        exit 1
                        ;;
                esac
                ;;
        esac
        ;;
    --help|-h|help)
        shift
        _help_verbose=0
        case "$(get_env)" in *dev*) _help_verbose=1 ;; esac
        for _a in "$@"; do
            case "$_a" in all|--verbose|-v) _help_verbose=1 ;; esac
        done
        show_help "$_help_verbose"
        ;;
    *)
        exec $PIXI run -e "$(get_env)" "$@"
        ;;
esac