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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
EXP_ROOT="${ROOT}/experiments/ablation"
RUN_ROOT="${EXP_ROOT}/background_runs"

BENCHMARK="bioagentbench"
VARIANT=""
RUN_ID=""
declare -a EXTRA_ARGS=()

usage() {
  cat <<'EOF'
Usage:
  launch_ablation_background.sh --benchmark {bioagentbench|labbench|all} --variant <variant> [extra args...]

Examples:
  launch_ablation_background.sh --benchmark bioagentbench --variant mcp_flat
  launch_ablation_background.sh --benchmark bioagentbench --variant mcp_metadata --task deseq
  launch_ablation_background.sh --benchmark labbench --variant biomanus --splits test
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --benchmark)
      BENCHMARK="$2"
      shift 2
      ;;
    --variant)
      VARIANT="$2"
      shift 2
      ;;
    --run-id)
      RUN_ID="$2"
      shift 2
      ;;
    --help|-h)
      usage
      exit 0
      ;;
    *)
      EXTRA_ARGS+=("$1")
      shift
      ;;
  esac
done

if [[ -z "${VARIANT}" ]]; then
  echo "Missing --variant" >&2
  usage >&2
  exit 2
fi

case "${BENCHMARK}" in
  bioagentbench)
    TARGET_SCRIPT="${EXP_ROOT}/scripts/run_bioagentbench_ablation.sh"
    ;;
  labbench)
    TARGET_SCRIPT="${EXP_ROOT}/scripts/run_labbench_ablation.sh"
    ;;
  all)
    TARGET_SCRIPT="${EXP_ROOT}/scripts/run_all_ablation.sh"
    ;;
  *)
    echo "Unsupported --benchmark: ${BENCHMARK}" >&2
    exit 2
    ;;
esac

timestamp="$(date -u +%Y%m%d_%H%M%S)"
if [[ -z "${RUN_ID}" ]]; then
  RUN_ID="${timestamp}_${BENCHMARK}_${VARIANT}"
fi

RUN_DIR="${RUN_ROOT}/${RUN_ID}"
mkdir -p "${RUN_DIR}"

CMD=("${TARGET_SCRIPT}" "--foreground" "--variant" "${VARIANT}" "${EXTRA_ARGS[@]}")

printf '%q ' "${CMD[@]}" > "${RUN_DIR}/command.sh"
printf '\n' >> "${RUN_DIR}/command.sh"
chmod +x "${RUN_DIR}/command.sh"
env | sort > "${RUN_DIR}/env.txt"

cat > "${RUN_DIR}/runner.sh" <<EOF
#!/usr/bin/env bash
set -euo pipefail
cd $(printf '%q' "${ROOT}")
echo "running" > $(printf '%q' "${RUN_DIR}/status.txt")
date -u +"%Y-%m-%dT%H:%M:%SZ START" >> $(printf '%q' "${RUN_DIR}/timeline.log")
if $(printf '%q ' "${CMD[@]}"); then
  rc=0
else
  rc=\$?
fi
echo "\${rc}" > $(printf '%q' "${RUN_DIR}/exit_code.txt")
if [[ "\${rc}" -eq 0 ]]; then
  echo "completed" > $(printf '%q' "${RUN_DIR}/status.txt")
else
  echo "failed" > $(printf '%q' "${RUN_DIR}/status.txt")
fi
date -u +"%Y-%m-%dT%H:%M:%SZ END rc=\${rc}" >> $(printf '%q' "${RUN_DIR}/timeline.log")
exit "\${rc}"
EOF
chmod +x "${RUN_DIR}/runner.sh"

setsid "${RUN_DIR}/runner.sh" > "${RUN_DIR}/stdout.log" 2>&1 < /dev/null &
launcher_pid=$!
echo "${launcher_pid}" > "${RUN_DIR}/pid"
echo "launched" > "${RUN_DIR}/status.txt"
date -u +"%Y-%m-%dT%H:%M:%SZ LAUNCHED pid=${launcher_pid}" > "${RUN_DIR}/timeline.log"

echo "Run directory: ${RUN_DIR}"
echo "PID: ${launcher_pid}"
echo "Status: ${RUN_DIR}/status.txt"
echo "Log: ${RUN_DIR}/stdout.log"