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

PROJECT_ROOT="/225040511/project/Biomni-ReAct"
PYTHON_BIN="${BIOMNI_REACT_PYTHON:-/225040511/miniconda3/envs/biomni_e1/bin/python}"
RUNNER="${PROJECT_ROOT}/scripts/run_labbench_dbqa.py"
OUTPUT_ROOT="${PROJECT_ROOT}/LAB-bench"
LOCK_FILE=""

DEV_SIZE="${LAB_BENCH_DEV_SIZE:-45}"
TEST_SIZE="${LAB_BENCH_TEST_SIZE:-315}"
SEED="${LAB_BENCH_SUBSET_SEED:-20260514}"
SHARD_COUNT="${LAB_BENCH_TEST_SHARD_COUNT:-10}"
SPLIT="${LAB_BENCH_SPLIT:-test}"

BACKGROUND=1
RESUME=1
while [[ $# -gt 0 ]]; do
  case "$1" in
    --foreground)
      BACKGROUND=0
      shift
      ;;
    --background)
      BACKGROUND=1
      shift
      ;;
    --fresh)
      RESUME=0
      shift
      ;;
    --resume)
      RESUME=1
      shift
      ;;
    --output-root)
      OUTPUT_ROOT="$2"
      shift 2
      ;;
    *)
      echo "Unknown argument: $1" >&2
      exit 2
      ;;
  esac
done

if [[ "${BACKGROUND}" -eq 1 ]]; then
  mkdir -p "${OUTPUT_ROOT}/logs"
  LOG="${OUTPUT_ROOT}/logs/dbqa_biomni_react_$(date -u +%Y%m%d_%H%M%S).log"
  SCRIPT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/$(basename "${BASH_SOURCE[0]}")"
  ARGS=(--foreground --output-root "${OUTPUT_ROOT}")
  if [[ "${RESUME}" -eq 1 ]]; then
    ARGS+=(--resume)
  else
    ARGS+=(--fresh)
  fi
  nohup "${SCRIPT}" "${ARGS[@]}" > "${LOG}" 2>&1 < /dev/null &
  echo "Started Biomni-ReAct LAB-Bench DbQA in background."
  echo "PID: $!"
  echo "Log: ${LOG}"
  echo "Results: ${OUTPUT_ROOT}/dbqa_results.jsonl"
  exit 0
fi

for ENV_FILE in "${PROJECT_ROOT}/.env" "/225040511/project/LAB-Bench/.env" "/225040511/project/.env"; do
  if [[ -f "${ENV_FILE}" ]]; then
    set -a
    # shellcheck disable=SC1090
    source "${ENV_FILE}"
    set +a
  fi
done

LOCK_FILE="${OUTPUT_ROOT}/dbqa_launcher.lock"
mkdir -p "${OUTPUT_ROOT}"
exec 9>"${LOCK_FILE}"
if ! flock -n 9; then
  echo "Another DbQA launcher is already running for ${OUTPUT_ROOT}. Lock: ${LOCK_FILE}" >&2
  exit 3
fi

if [[ -z "${BIOMNI_REACT_API_KEY:-}" && -z "${DEEPSEEK_API_KEY:-}" && -z "${BIOMNI_CUSTOM_API_KEY:-}" && -z "${OPENAI_API_KEY:-}" ]]; then
  echo "No LLM API key found. Set DEEPSEEK_API_KEY, BIOMNI_REACT_API_KEY, BIOMNI_CUSTOM_API_KEY, or OPENAI_API_KEY." >&2
  exit 2
fi

export BIOMNI_REACT_MODEL="${BIOMNI_REACT_MODEL:-${DEEPSEEK_MODEL_NAME:-deepseek-chat}}"
export BIOMNI_REACT_BASE_URL="${BIOMNI_REACT_BASE_URL:-${DEEPSEEK_BASE_URL:-https://api.deepseek.com/v1}}"

mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/shard_logs"
RESULT_FILE="${OUTPUT_ROOT}/dbqa_results.jsonl"
REASONING_LOG="${OUTPUT_ROOT}/dbqa_reasoning.log"
if [[ "${RESUME}" -eq 0 ]]; then
  : > "${RESULT_FILE}"
  : > "${REASONING_LOG}"
fi
touch "${RESULT_FILE}" "${REASONING_LOG}"

declare -a PIDS=()
for (( shard_index=0; shard_index<SHARD_COUNT; shard_index++ )); do
  SHARD_LOG="${OUTPUT_ROOT}/shard_logs/dbqa_shard$(printf '%02d' "$((shard_index + 1))")of$(printf '%02d' "${SHARD_COUNT}").log"
  "${PYTHON_BIN}" "${RUNNER}" \
    --split "${SPLIT}" \
    --dev-size "${DEV_SIZE}" \
    --test-size "${TEST_SIZE}" \
    --seed "${SEED}" \
    --shard-index "${shard_index}" \
    --shard-count "${SHARD_COUNT}" \
    --output-root "${OUTPUT_ROOT}" \
    --result-file "${RESULT_FILE}" \
    --reasoning-log "${REASONING_LOG}" \
    --skip-existing-results \
    > "${SHARD_LOG}" 2>&1 &
  PIDS+=("$!")
  echo "[launcher] started shard $((shard_index + 1))/${SHARD_COUNT} pid=${PIDS[-1]} log=${SHARD_LOG}"
done

FAILED=0
for pid in "${PIDS[@]}"; do
  if ! wait "${pid}"; then
    FAILED=1
  fi
done

COUNT=$(wc -l < "${RESULT_FILE}" || true)
echo "[launcher] dbqa result lines=${COUNT}; file=${RESULT_FILE}"
if [[ "${FAILED}" -ne 0 ]]; then
  echo "[launcher] one or more shards failed" >&2
  exit 1
fi