#!/bin/bash
# Per-node stage runner for the Russell 3000 extension (2026-10-27).
#   run_shard.sh <node> <stage> [args]
# stages: download_self | build_db | parse | embed | export | pipeline (build_db→parse→embed→export)
# Each stage writes data/logs/<stage>.log and data/logs/stage.<stage>.{running,done,failed}.
set -u
NODE="$1"; STAGE="$2"; shift 2
ROOT="$HOME/UQO_RESEARCH_DB"
cd "$ROOT" || exit 1
export PATH="$HOME/.local/bin:/opt/homebrew/bin:$PATH"
export PYTHONUNBUFFERED=1
PY="$ROOT/.venv/bin/python"
LOGS="$ROOT/data/logs"; mkdir -p "$LOGS" "$ROOT/data/metadata/shards"
RATE="${RATE:-0.9}"; WORKERS="${WORKERS:-8}"; DL_WORKERS="${DL_WORKERS:-6}"

mark() { rm -f "$LOGS/stage.$1.running" "$LOGS/stage.$1.done" "$LOGS/stage.$1.failed"; touch "$LOGS/stage.$1.$2"; }

run_stage() {  # name, command...
  local name="$1"; shift
  mark "$name" running
  echo "=== $(date '+%F %T') [$NODE] stage $name start: $*" >> "$LOGS/$name.log"
  "$@" >> "$LOGS/$name.log" 2>&1
  local rc=$?
  if [ $rc -eq 0 ]; then mark "$name" done; else mark "$name" failed; fi
  echo "=== $(date '+%F %T') [$NODE] stage $name rc=$rc" >> "$LOGS/$name.log"
  return $rc
}

case "$STAGE" in
  download_self)
    LIST="data/metadata/shards/$NODE.self.txt"
    if [ ! -s "$LIST" ]; then mark download_self done; exit 0; fi
    KNOWN=""; [ -s data/metadata/download_log.json ] && KNOWN="--known-logs data/metadata/download_log.json"
    run_stage download_self "$PY" scripts/download_filings.py --ciks-file "$LIST" --rate "$RATE" \
        --workers "$DL_WORKERS" --log data/metadata/download_log.self.jsonl $KNOWN ;;
  build_db)
    run_stage build_db "$PY" scripts/build_database.py --shard-node "$NODE" --no-parquet ;;
  parse)
    run_stage parse "$PY" scripts/parse_filings.py --max-workers "$WORKERS" --batch-size 800 --no-json ;;
  embed)
    run_stage embed "$PY" scripts/build_embeddings.py --batch-size 256 ;;
  export)
    run_stage export "$PY" scripts/export_shard.py ;;
  pipeline)
    mark pipeline running
    "$0" "$NODE" build_db && "$0" "$NODE" parse && "$0" "$NODE" embed && "$0" "$NODE" export
    rc=$?; if [ $rc -eq 0 ]; then mark pipeline done; else mark pipeline failed; fi; exit $rc ;;
  pipeline_master)
    mark pipeline running
    "$0" "$NODE" build_db && "$0" "$NODE" parse && "$0" "$NODE" embed
    rc=$?; if [ $rc -eq 0 ]; then mark pipeline done; else mark pipeline failed; fi; exit $rc ;;
  *) echo "unknown stage $STAGE"; exit 2 ;;
esac