#!/usr/bin/env bash
set -euo pipefail

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd -P)"
cd "$repo_root"

if [[ "${PARTY2D_APPROVE_MODEL_FIT:-}" != "YES" ]]; then
  echo "Refusing to start a Stan fit without explicit approval." >&2
  echo "After approval, rerun with PARTY2D_APPROVE_MODEL_FIT=YES." >&2
  exit 64
fi

run_dir="${1:-$repo_root/_local/validation/blocked_party}"
mkdir -p "$run_dir" "$repo_root/_local/tmp/blocked_validation"

if pgrep -af 'stan_model_2dim_v6.*sample|src/julia/01_run_model.jl' | grep -v "$$" >/dev/null; then
  echo "Refusing to start: another party2d model process appears to be running." >&2
  pgrep -af 'stan_model_2dim_v6.*sample|src/julia/01_run_model.jl' >&2 || true
  exit 1
fi

if [[ -f "$run_dir/model_fit.complete" ]]; then
  echo "Blocked validation already marked complete: $run_dir/model_fit.complete" >&2
  exit 1
fi

required_kb=$((80 * 1024 * 1024))
available_kb="$(df -Pk "$run_dir" | awk 'NR==2 {print $4}')"
if (( available_kb < required_kb )); then
  echo "Refusing to start: less than 80 GiB free on the run volume." >&2
  exit 1
fi

julia --project=. validation/prepare_blocked_validation.jl "$run_dir"
sha256sum "$run_dir"/text_data.csv "$run_dir"/expert.csv "$run_dir"/lr_data.csv \
  "$run_dir"/expert_test.csv "$run_dir"/lr_data_test.csv \
  > "$run_dir/input_sha256sums.txt"

export CMDSTAN_HOME="${CMDSTAN_HOME:-/opt/agent-tools/cmdstan-2.39.0}"
export TMPDIR="$repo_root/_local/tmp/blocked_validation"
export JULIA_NUM_THREADS="${JULIA_NUM_THREADS:-4}"
export STAN_NUM_THREADS=1
export STAN_REFRESH="${STAN_REFRESH:-50}"
export PARTY2D_NUM_CHAINS="${PARTY2D_NUM_CHAINS:-4}"
export PARTY2D_NUM_WARMUP="${PARTY2D_NUM_WARMUP:-1000}"
export PARTY2D_NUM_SAMPLES="${PARTY2D_NUM_SAMPLES:-1000}"

if [[ ! -x "$CMDSTAN_HOME/bin/stanc" ]]; then
  echo "Refusing to start: no executable stanc compiler at $CMDSTAN_HOME/bin/stanc" >&2
  exit 1
fi
"$CMDSTAN_HOME/bin/stanc" --version > "$run_dir/stanc_version.txt"

date --iso-8601=seconds > "$run_dir/model_fit.started"
echo "$$" > "$run_dir/model_fit.pid"
echo "CMDSTAN_HOME=$CMDSTAN_HOME" > "$run_dir/model_fit.environment"
echo "JULIA_NUM_THREADS=$JULIA_NUM_THREADS" >> "$run_dir/model_fit.environment"
echo "STAN_NUM_THREADS=$STAN_NUM_THREADS" >> "$run_dir/model_fit.environment"
echo "PARTY2D_NUM_CHAINS=$PARTY2D_NUM_CHAINS" >> "$run_dir/model_fit.environment"
echo "PARTY2D_NUM_WARMUP=$PARTY2D_NUM_WARMUP" >> "$run_dir/model_fit.environment"
echo "PARTY2D_NUM_SAMPLES=$PARTY2D_NUM_SAMPLES" >> "$run_dir/model_fit.environment"

set +e
nice -n 10 julia --project=. src/julia/01_run_model.jl --data-dir "$run_dir" \
  2>&1 | tee "$run_dir/model_fit.log"
status="${PIPESTATUS[0]}"
set -e

echo "$status" > "$run_dir/model_fit.exit_code"
if [[ "$status" -eq 0 ]]; then
  rm -f "$run_dir/model_fit.failed"
  date --iso-8601=seconds > "$run_dir/model_fit.complete"
else
  date --iso-8601=seconds > "$run_dir/model_fit.failed"
fi
exit "$status"