#!/usr/bin/env bash
# Worker B has already claimed chunk 8, so worker A must NOT walk on to it after
# finishing chunk 4 - it would duplicate ~1h of reranking and finish later than B.
#
# Plan: let A finish chunk 4, stop A and its supervisor, wait for B to land chunk 8,
# then run ONE decision-only pass over all 9 chunks.
set -uo pipefail
cd "$(dirname "$0")/.."
ROOT=$(pwd); LOGS="$ROOT/logs"; CK="$ROOT/work/results_A/test_chunks"
PY=/home/user4/Jyo/DAI/er_pipeline/.venv/bin/python
WD="$LOGS/finalize.log"
say() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$WD"; }

say "=== waiting for chunk 4 (A) and chunk 8 (B) ==="
while true; do
  n=$(ls "$CK"/chunk_*.parquet 2>/dev/null | wc -l)
  have4=0; [ -f "$CK/chunk_000800000.parquet" ] && have4=1
  # once chunk 4 exists, stop A so it cannot start duplicating chunk 8
  if [ "$have4" = 1 ] && [ ! -f "$LOGS/.a_stopped" ]; then
    for s in $(pgrep -f "[s]upervise.sh test"); do say "stopping supervisor $s"; kill -9 "$s"; done
    for p in $(pgrep -f "[r]un_test.py"); do
      ps -p "$p" -o cmd --no-headers 2>/dev/null | grep -q "model-3/work/results_A" || continue
      tr '\0' '\n' < "/proc/$p/environ" 2>/dev/null | grep -q "M3_CHUNK_LO" && continue
      say "stopping worker A pid $p (chunk 4 done; chunk 8 belongs to B)"; kill -9 "$p"
    done
    touch "$LOGS/.a_stopped"
  fi
  [ "$n" -ge 9 ] && { say "all 9 chunks present"; break; }
  sleep 60
done

say "running the global decision pass (conflict resolution + selection + write + validate)"
env CUDA_VISIBLE_DEVICES=1 HF_HOME=/scratch/user4/mtp/cache/huggingface \
  PYTHONPATH="$ROOT/src" TOKENIZERS_PARALLELISM=false \
  "$PY" "$ROOT/src/run_test.py" --gpu 0 --artifacts "$ROOT/work/results_A" \
  >> "$LOGS/run_test_final.log" 2>&1
rc=$?
say "decision pass exited rc=$rc"
if [ -f "$ROOT/output/matching_results.tsv" ]; then
  say "matching_results.tsv: $(wc -l < "$ROOT/output/matching_results.tsv") lines"
  say "candidate_pairs.tsv : $(wc -l < "$ROOT/output/candidate_pairs.tsv") lines"
fi
say "=== finalize done ==="