#!/usr/bin/env bash
# Gather per-seed HEALTH_DIAG metrics from a multi-seed Argo run and pipe
# them through aggregate-multi-seed-metrics.py for cross-seed statistics.
#
# Plan 5 Task 1B.1.
#
# Usage:
#   ./scripts/gather-multi-seed-metrics.sh <tag> [output.json]
#
# Where <tag> matches the foxhunt-tag label set via `argo-train.sh --tag`.
# The script:
#   1. Pulls all logs for workflows labelled foxhunt-tag=<tag> into a
#      flat text file (/tmp/all-logs-<tag>.txt).
#   2. Invokes aggregate-multi-seed-metrics.py to produce a single JSON
#      report with per-(epoch, metric) mean/std/median/min/max across
#      seeds.
#
# Defaults: --multi-seed and --folds are read from the MULTI_SEED / FOLDS
# environment vars (default 5 / 6). Override via env or by editing here.
set -euo pipefail

TAG="${1:?Usage: gather-multi-seed-metrics.sh <tag> [output.json]}"
OUT="${2:-/tmp/multi-seed-${TAG}.json}"
LOG_FILE="/tmp/all-logs-${TAG}.txt"

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"

if ! command -v argo >/dev/null 2>&1; then
  echo "ERROR: argo CLI not found in PATH" >&2
  exit 1
fi

if ! command -v python3 >/dev/null 2>&1; then
  echo "ERROR: python3 not found in PATH" >&2
  exit 1
fi

echo "=== Pulling logs for workflows tagged foxhunt-tag=${TAG} ==="
# `argo logs @latest` is a single-workflow tail; for multi-workflow
# selection we list workflows by label, then concatenate their logs.
WORKFLOWS=$(argo list -n foxhunt --selector "foxhunt-tag=${TAG}" -o name 2>/dev/null || true)
if [[ -z "${WORKFLOWS}" ]]; then
  echo "ERROR: no workflows found with foxhunt-tag=${TAG}" >&2
  exit 1
fi

: > "${LOG_FILE}"
for wf in ${WORKFLOWS}; do
  echo "  -> $wf"
  argo logs -n foxhunt --no-color "$wf" >> "${LOG_FILE}" 2>&1 || {
    echo "  WARN: failed to fetch logs for $wf" >&2
  }
done

LOG_LINES=$(wc -l < "${LOG_FILE}")
echo "Collected ${LOG_LINES} log lines → ${LOG_FILE}"

echo "=== Aggregating per-(epoch, metric) statistics ==="
python3 "${REPO_ROOT}/scripts/aggregate-multi-seed-metrics.py" \
  --input "${LOG_FILE}" \
  --multi-seed "${MULTI_SEED:-5}" \
  --folds      "${FOLDS:-6}" \
  --tag        "${TAG}" \
  --output     "${OUT}"

echo "Aggregated → ${OUT}"