#!/bin/bash
# Federated Grid Fault Detection — Phase A orchestrator (BRING-UP ONLY).
#
# LoomAI's "Run" button invokes this with the slice name as $1 (and also exports
# the weave.json args as env vars). This weave provisions the slice and EXITS when
# it is READY — it does not run federated learning and does not poll in a loop.
# Run the FL from federated_learning.ipynb once the slice is up.
#
# Manual use:
#   ./weave.sh <slice-name>            # start (default)
#   ACTION=status ./weave.sh <slice>   # one-shot health check
#   ACTION=stop   ./weave.sh <slice>   # delete the slice
# Teardown is normally done via weave_cleanup.sh (the LoomAI cleanup script).
set -uo pipefail

# Slice name: env SLICE_NAME wins, else $1, else default (matches LoomAI convention).
SLICE_NAME="${SLICE_NAME:-${1:-fl-grid}}"
SLICE_NAME=$(echo "$SLICE_NAME" | sed 's/[^a-zA-Z0-9-]/-/g; s/--*/-/g; s/^-//; s/-$//')
ACTION="${ACTION:-start}"

# Provisioning args (weave.md §6.1) — exported for experiment.py.
export NUM_CLIENTS="${NUM_CLIENTS:-4}"
export USE_GPU="${USE_GPU:-true}"
export GPU_MODEL="${GPU_MODEL:-auto}"
export AGG_GPU="${AGG_GPU:-false}"
export SITE_MAP="${SITE_MAP:-}"
export EXCLUDE_SITES="${EXCLUDE_SITES:-}"
export CUDA_TAG="${CUDA_TAG:-cu121}"
export KEEP_SLICE="${KEEP_SLICE:-true}"

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"

# If the run is stopped mid-provision (SIGTERM/SIGINT), abort and tear down the
# partial slice. After READY the process has already exited, so the slice persists.
cleanup() {
    echo "### PROGRESS: Stop requested during bring-up — deleting partial slice..."
    python3 "$SCRIPT_DIR/experiment.py" stop "$SLICE_NAME" 2>&1 || true
    exit 0
}
trap cleanup SIGTERM SIGINT

if [ "$ACTION" = "start" ]; then
    echo "### PROGRESS: === Federated Grid Fault Detection — bring-up ==="
    echo "### PROGRESS: Slice:    $SLICE_NAME"
    echo "### PROGRESS: Clients:  $NUM_CLIENTS"
    if [ "$USE_GPU" = "true" ]; then
        echo "### PROGRESS: Profile:  GPU (model=$GPU_MODEL, agg_gpu=$AGG_GPU)"
    else
        echo "### PROGRESS: Profile:  CPU-edge"
    fi
    [ -n "$SITE_MAP" ] && echo "### PROGRESS: SiteMap:  $SITE_MAP"
    echo ""
fi

python3 "$SCRIPT_DIR/experiment.py" "$ACTION" "$SLICE_NAME"
RC=$?

if [ "$ACTION" = "start" ] && [ $RC -eq 0 ]; then
    echo ""
    echo "### PROGRESS: Next: open federated_learning.ipynb and run the FL against '$SLICE_NAME'."
    echo "### PROGRESS: The slice stays up. Tear down with: ./weave_cleanup.sh $SLICE_NAME"
fi
exit $RC
