#!/usr/bin/env bash
# Farm multi-build pieces from jobs/pieces.yaml (or PIECES_FILE) onto job-bus.
set -euo pipefail
export PATH="${HOME}/bin:/opt/homebrew/bin:/usr/local/bin:${PATH}"
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
cd "$ROOT"

if [[ -f "${HOME}/.config/nats/job-bus.env" ]]; then
  set -a; # shellcheck disable=SC1090
  source "${HOME}/.config/nats/job-bus.env"; set +a
elif [[ -f /etc/nats/job-bus.env ]]; then
  set -a; source /etc/nats/job-bus.env; set +a
else
  echo "missing job-bus.env" >&2; exit 1
fi

PIECES_FILE="${PIECES_FILE:-$ROOT/jobs/pieces.yaml}"
[[ -f "$PIECES_FILE" ]] || { echo "missing $PIECES_FILE (copy from job-bus template)" >&2; exit 1; }

PY="${JOB_BUS_PYTHON:-}"
if [[ -z "$PY" ]]; then
  if [[ -x "$HOME/job-bus/venv/bin/python" ]]; then PY="$HOME/job-bus/venv/bin/python"
  elif [[ -x /opt/job-bus/venv/bin/python ]]; then PY=/opt/job-bus/venv/bin/python
  else PY=python3
  fi
fi

WATCH=0
CLEAN_AFTER=0
PROGRESS=0
while [[ $# -gt 0 ]]; do
  case "$1" in
    --watch) WATCH=1; shift ;;
    --progress) PROGRESS=1; shift ;;
    --cleanup-after) CLEAN_AFTER=1; shift ;;
    -h|--help)
      echo "usage: scripts/jobbus-farm.sh [--watch] [--progress] [--cleanup-after]"
      exit 0 ;;
    *) echo "unknown: $1" >&2; exit 1 ;;
  esac
done

# Ensure PyYAML available in venv lightly
"$PY" -c 'import yaml' 2>/dev/null || "$PY" -m pip install -q pyyaml

MAP=$("$PY" - "$PIECES_FILE" "$ROOT" <<'PY'
import json, os, sys, datetime
from pathlib import Path
try:
    import yaml
except ImportError:
    raise SystemExit("PyYAML required")

path, root = sys.argv[1], sys.argv[2]
doc = yaml.safe_load(open(path))
defaults = doc.get("defaults") or {}
repo = Path(root).name
date = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
batch = (doc.get("batch") or f"{repo}-{date}").replace("{{REPO_NAME}}", repo).replace("{{DATE}}", date)
pieces = doc.get("pieces") or []
out_dir = Path(root) / "jobs" / "submitted" / batch
out_dir.mkdir(parents=True, exist_ok=True)
ids = []
for i, p in enumerate(pieces, 1):
    pid = str(p.get("id") or f"piece-{i:02d}")
    kind = p.get("kind") or defaults.get("kind") or "cpu"
    requires = p.get("requires") if "requires" in p else defaults.get("requires") or []
    if requires is None:
        requires = []
    cmd = p.get("cmd")
    if not cmd:
        raise SystemExit(f"piece {pid} missing cmd")
    job_id = f"job-{batch}-{pid}"
    body = {
        "id": job_id,
        "kind": kind,
        "requires": list(requires),
        "cmd": list(cmd),
        "timeout_sec": int(p.get("timeout_sec") or defaults.get("timeout_sec") or 3600),
        "cleanup": p.get("cleanup") or defaults.get("cleanup") or "none",
        "meta": {"batch": batch, "piece": pid, "repo": repo, **(p.get("meta") or {})},
        "ts": datetime.datetime.now(datetime.timezone.utc).isoformat().replace("+00:00", "Z"),
    }
    if p.get("cwd"):
        body["cwd"] = p["cwd"]
    else:
        # default cwd: repo root — only works if worker can see this path (shared mount).
        # Until SSHFS is set, prefer pieces that use absolute worker paths or staged inputs.
        body["cwd"] = root
    jf = out_dir / f"{pid}.json"
    jf.write_text(json.dumps(body, indent=2) + "\n")
    ids.append({"id": job_id, "kind": kind, "file": str(jf), "piece": pid})
print(json.dumps({"batch": batch, "jobs": ids, "count": len(ids)}))
PY
)

BATCH=$(echo "$MAP" | "$PY" -c 'import json,sys; print(json.load(sys.stdin)["batch"])')
COUNT=$(echo "$MAP" | "$PY" -c 'import json,sys; print(json.load(sys.stdin)["count"])')
echo "batch=$BATCH pieces=$COUNT"

echo "$MAP" | "$PY" -c 'import json,sys; d=json.load(sys.stdin); print("\n".join(j["file"] for j in d["jobs"]))' | while read -r jf; do
  submit-job --json "$jf"
done

echo "submitted $COUNT jobs for batch $BATCH"
echo "track: job-progress --batch $BATCH --expect $COUNT"
echo "clean: job-cleanup --batch $BATCH"

if [[ "$PROGRESS" -eq 1 ]]; then
  job-progress --batch "$BATCH" --expect "$COUNT" || true
fi

if [[ "$WATCH" -eq 1 ]]; then
  # simple wait via status subscription count*2 (running+done) approx using progress expect
  job-progress --batch "$BATCH" --expect "$COUNT" || true
fi

if [[ "$CLEAN_AFTER" -eq 1 ]]; then
  job-cleanup --batch "$BATCH" || true
fi