#!/usr/bin/env bash
# Run the corpus pipeline after the download completes.
# Portable across jupyter-pod and kls-headnode.
#
#   bash scripts/next.sh              # check prerequisites only, run nothing
#   bash scripts/next.sh --run        # run schema -> corpus -> fertility
set -uo pipefail

REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
ANURAG="$(dirname "$REPO")"
ROOT="${VOICERAG_ROOT:-$ANURAG/voicerag_data}"
RUN=0
[ "${1:-}" = "--run" ] && RUN=1

cd "$REPO"
mkdir -p logs

echo "════════ CONTEXT ════════"
echo "  host : $(hostname -s)"
echo "  repo : $REPO"
echo "  root : $ROOT"

# ---------------------------------------------------------------- 1. data
echo
echo "════════ 1. DATASET ════════"
HUB="$ROOT/hf_cache/hub"
if [ -d "$HUB" ]; then
  echo "  hf_cache: $(du -sh "$ROOT/hf_cache" 2>/dev/null | cut -f1)"
  for r in datasets--ai4bharat--MSMARCO-XI models--BAAI--bge-m3 \
           models--BAAI--bge-reranker-v2-m3 models--sentence-transformers--all-MiniLM-L6-v2; do
    [ -d "$HUB/$r" ] && echo "    ok      $r" || echo "    MISSING $r"
  done
  NPQ=$(find "$HUB" -name '*.parquet' 2>/dev/null | wc -l)
  echo "  parquet files: $NPQ"
  # HF caches use RELATIVE symlinks (snapshots -> ../../blobs), so a move is safe.
  # Verify one actually resolves after any relocation.
  SAMPLE=$(find "$HUB" -name '*.parquet' 2>/dev/null | head -1)
  if [ -n "$SAMPLE" ]; then
    if [ -r "$SAMPLE" ] && [ -s "$SAMPLE" ]; then
      echo "  symlinks resolve: ok  ($(du -Lh "$SAMPLE" 2>/dev/null | cut -f1) sample)"
    else
      echo "  !! symlinks BROKEN after move — re-run 01_download.py to repair"
      exit 1
    fi
  fi
else
  echo "  !! no hf_cache at $ROOT"
  echo "     If the download ran elsewhere: bash scripts/relocate.sh"
  exit 1
fi

# ---------------------------------------------------------------- 2. python
echo
echo "════════ 2. PYTHON DEPENDENCIES ════════"
MISSING=$(python3 - <<'PY'
need = {"polars":"schema+corpus", "pyarrow":"schema", "numpy":"all",
        "transformers":"fertility", "matplotlib":"fertility plot"}
missing = []
for m, why in need.items():
    try:
        mod = __import__(m)
        print(f"  ok      {m:14s} {getattr(mod,'__version__','?'):12s} ({why})")
    except Exception:
        print(f"  MISSING {m:14s} {'':12s} ({why})")
        missing.append(m)
print("MISSING:" + ",".join(missing))
PY
)
echo "$MISSING" | grep -v '^MISSING:'
MISS=$(echo "$MISSING" | sed -n 's/^MISSING://p')
if [ -n "$MISS" ]; then
  echo
  echo "  !! Missing: $MISS"
  echo "     These exist on jupyter-pod. Either run this there, or create an"
  echo "     isolated venv here (nothing system-wide is touched):"
  echo
  echo "       python3 -m venv --system-site-packages $ANURAG/.venv"
  echo "       source $ANURAG/.venv/bin/activate"
  echo "       pip install ${MISS//,/ }"
  echo
  [ "$RUN" -eq 1 ] && { echo "  aborting."; exit 1; }
fi

# ---------------------------------------------------------------- 3. status
echo
echo "════════ 3. PIPELINE STATUS ════════"
step() { [ -e "$2" ] && echo "  done    $1  ($(du -sh "$2" 2>/dev/null | cut -f1))" \
                     || echo "  pending $1"; }
step "schema_report.json " "$ROOT/data/schema_report.json"
step "passages.parquet   " "$ROOT/data/passages.parquet"
step "pseudo_docs.parquet" "$ROOT/data/pseudo_docs.parquet"
step "fertility.json     " "$ROOT/results/fertility.json"

if [ "$RUN" -eq 0 ]; then
  cat <<EOF

════════ TO RUN ════════
  export VOICERAG_ROOT="$ROOT"
  bash scripts/next.sh --run

  Or step by step (each writes only under \$VOICERAG_ROOT):
    python3 scripts/02_inspect_schema.py --root "\$VOICERAG_ROOT"
    python3 src/pseudo_docs.py           --root "\$VOICERAG_ROOT" --strategy url
    python3 src/fertility.py             --root "\$VOICERAG_ROOT"
EOF
  exit 0
fi

export VOICERAG_ROOT="$ROOT"

# ---------------------------------------------------------------- run
echo
echo "════════ RUNNING ════════"

echo
echo "-- step 1/3: schema inspection (fast) --"
python3 scripts/02_inspect_schema.py --root "$ROOT" 2>&1 | tee logs/schema.log | tail -45
[ -f "$ROOT/data/schema_report.json" ] || { echo "  !! schema report not written; see logs/schema.log"; exit 1; }

echo
echo "-- step 2/3: fertility (needs only the schema report; ~5 min) --"
python3 src/fertility.py --root "$ROOT" 2>&1 | tee logs/fertility.log | tail -30

echo
echo "-- step 3/3: corpus build (memory-heavy; backgrounded) --"
nohup python3 src/pseudo_docs.py --root "$ROOT" --strategy url --target-words 3000 \
      > logs/corpus.log 2>&1 &
echo "  started pid $! -> logs/corpus.log"
echo "  watch: tail -f $REPO/logs/corpus.log"

cat <<EOF

════════ SEND ME ════════
  $ROOT/data/schema_report.json
  $ROOT/results/fertility.json
  $REPO/logs/corpus.log        (once it finishes)
EOF