#!/usr/bin/env bash
set -u

RUN_ID="${RUN_ID:-cloning_50shards_$(date -u +%Y%m%d_%H%M%S)}"
PROJECT_ROOT="/225040511/project"
PYTHON_BIN="${LAB_BENCH_RUNNER_PYTHON:-/225040511/miniconda3/envs/biomni_e1/bin/python}"
LOG_ROOT="${PROJECT_ROOT}/cloningscenarios_logs/${RUN_ID}"
SHARD_COUNT="${CLONING_SHARD_COUNT:-50}"
SHARD_CONCURRENCY="${CLONING_SHARD_CONCURRENCY:-5}"
DEV_SIZE="${CLONING_DEV_SIZE:-0}"
TEST_SIZE="${CLONING_TEST_SIZE:-33}"
SEED="${CLONING_SEED:-20260514}"

mkdir -p "${LOG_ROOT}"

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

export DEEPSEEK_BASE_URL="${DEEPSEEK_BASE_URL:-https://api.deepseek.com/v1}"
export DEEPSEEK_MODEL_NAME="${DEEPSEEK_MODEL_NAME:-deepseek-chat}"
export REACT_CODE_API_KEY="${REACT_CODE_API_KEY:-${DEEPSEEK_API_KEY:-${BIOMNI_CUSTOM_API_KEY:-}}}"
export REACT_CODE_BASE_URL="${REACT_CODE_BASE_URL:-${DEEPSEEK_BASE_URL}}"
export REACT_CODE_MODEL="${REACT_CODE_MODEL:-${DEEPSEEK_MODEL_NAME}}"

wait_limited() {
  local -n pid_array="$1"
  if [[ "${#pid_array[@]}" -lt "${SHARD_CONCURRENCY}" ]]; then
    return
  fi
  local pid="${pid_array[0]}"
  wait "${pid}" || true
  pid_array=("${pid_array[@]:1}")
}

count_rows() {
  local file="$1"
  if [[ -f "${file}" ]]; then
    wc -l < "${file}"
  else
    echo 0
  fi
}

run_react_code_bioagent() {
  local runner="${PROJECT_ROOT}/React+code_labbench/run_labbench_react_code.py"
  local out="${PROJECT_ROOT}/react_code_bioagent_deepseek/labbench_runs/react_code_cloningscenarios_${RUN_ID}"
  local result="${out}/cloningscenarios_results.jsonl"
  local reasoning="${out}/cloningscenarios_reasoning.log"
  mkdir -p "${out}/shard_logs"
  touch "${result}" "${reasoning}"
  echo "react_code_bioagent output=${out}"
  echo "react_code_bioagent before_rows=$(count_rows "${result}")"
  local pids=()
  for ((shard=0; shard<SHARD_COUNT; shard++)); do
    local label
    label="$(printf "react_code_bioagent_cloningscenarios_shard%02dof%02d" "$((shard + 1))" "${SHARD_COUNT}")"
    "${PYTHON_BIN}" "${runner}" \
      --eval CloningScenarios \
      --split test \
      --dev-size "${DEV_SIZE}" \
      --test-size "${TEST_SIZE}" \
      --seed "${SEED}" \
      --shard-index "${shard}" \
      --shard-count "${SHARD_COUNT}" \
      --output-root "${out}" \
      --result-file "${result}" \
      --reasoning-log "${reasoning}" \
      --skip-existing-results \
      > "${out}/shard_logs/${label}.log" 2>&1 &
    pids+=("$!")
    wait_limited pids
  done
  for pid in "${pids[@]}"; do wait "${pid}" || true; done
  echo "react_code_bioagent after_rows=$(count_rows "${result}")"
}

run_base_deepseek() {
  local runner="${PROJECT_ROOT}/react_code_bioagent_deepseek/react_code_bioagent/labbench_runner.py"
  local out="${PROJECT_ROOT}/base_llm_labbench/results/cloningscenarios_base_deepseek_${RUN_ID}"
  local result="${out}/cloningscenarios_results.jsonl"
  local reasoning="${out}/cloningscenarios_reasoning.log"
  mkdir -p "${out}/shard_logs"
  touch "${result}" "${reasoning}"
  echo "base_deepseek output=${out}"
  echo "base_deepseek before_rows=$(count_rows "${result}")"
  local pids=()
  for ((shard=0; shard<SHARD_COUNT; shard++)); do
    local label
    label="$(printf "base_deepseek_cloningscenarios_shard%02dof%02d" "$((shard + 1))" "${SHARD_COUNT}")"
    "${PYTHON_BIN}" "${runner}" \
      --evals CloningScenarios \
      --splits test \
      --dev-size "${DEV_SIZE}" \
      --test-size "${TEST_SIZE}" \
      --seed "${SEED}" \
      --n-threads 1 \
      --shard-index "${shard}" \
      --shard-count "${SHARD_COUNT}" \
      --output-root "${out}" \
      --compact-results-path "${result}" \
      --reasoning-log-path "${reasoning}" \
      --skip-existing-results \
      > "${out}/shard_logs/${label}.log" 2>&1 &
    pids+=("$!")
    wait_limited pids
  done
  for pid in "${pids[@]}"; do wait "${pid}" || true; done
  echo "base_deepseek after_rows=$(count_rows "${result}")"
}

echo "RUN_ID=${RUN_ID}"
echo "SHARD_COUNT=${SHARD_COUNT}"
echo "SHARD_CONCURRENCY=${SHARD_CONCURRENCY}"

run_react_code_bioagent > "${LOG_ROOT}/react_code_bioagent_deepseek.log" 2>&1 &
pid_react=$!
run_base_deepseek > "${LOG_ROOT}/base_deepseek.log" 2>&1 &
pid_base=$!

echo "react_code_bioagent_pid=${pid_react}"
echo "base_deepseek_pid=${pid_base}"
echo "log_root=${LOG_ROOT}"

wait "${pid_react}" || true
wait "${pid_base}" || true
echo "finished RUN_ID=${RUN_ID}"