#!/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"
RUNNER="${ROOT}/experiments/lab_bench/scripts/run_labbench_with_hypobioos.py"
PYTHON_BIN="${BIOMANUS_PYTHON:-/225040511/miniconda3/envs/biomni_e1/bin/python}"
GRAPH_BUILDER="${ROOT}/build_generated_mcp_graph.py"
GRAPH_SANITIZER="${ROOT}/experiments/lab_bench/scripts/sanitize_labbench_mcp_graph.py"

BACKGROUND=0
VARIANT="all"
DEBUG=0
DEV_SIZE="${BIOMANUS_LABBENCH_DEV_SIZE:-45}"
TEST_SIZE="${BIOMANUS_LABBENCH_TEST_SIZE:-315}"
SHARD_COUNT="${BIOMANUS_LABBENCH_SHARD_COUNT:-10}"
SPLITS_TEXT="${BIOMANUS_LABBENCH_SPLITS:-test}"

while [[ $# -gt 0 ]]; do
  case "$1" in
    --background) BACKGROUND=1; shift ;;
    --foreground) BACKGROUND=0; shift ;;
    --variant) VARIANT="$2"; shift 2 ;;
    --debug) DEBUG=1; shift ;;
    --dev-size) DEV_SIZE="$2"; shift 2 ;;
    --test-size) TEST_SIZE="$2"; shift 2 ;;
    --shard-count) SHARD_COUNT="$2"; shift 2 ;;
    --splits) SPLITS_TEXT="$2"; shift 2 ;;
    *) echo "Unknown argument: $1" >&2; exit 2 ;;
  esac
done

if [[ "${BACKGROUND}" -eq 1 ]]; then
  mkdir -p "${EXP_ROOT}/logs"
  LOG="${EXP_ROOT}/logs/labbench_ablation_$(date -u +%Y%m%d_%H%M%S).log"
  SCRIPT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/$(basename "${BASH_SOURCE[0]}")"
  CMD=(
    "${SCRIPT}" "--foreground"
    "--variant" "${VARIANT}"
    "--dev-size" "${DEV_SIZE}"
    "--test-size" "${TEST_SIZE}"
    "--shard-count" "${SHARD_COUNT}"
    "--splits" "${SPLITS_TEXT}"
  )
  if [[ "${DEBUG}" -eq 1 ]]; then CMD+=("--debug"); fi
  nohup "${CMD[@]}" > "${LOG}" 2>&1 < /dev/null &
  echo "Started LAB-Bench ablation in background."
  echo "PID: $!"
  echo "Log: ${LOG}"
  exit 0
fi

if [[ -f "${ROOT}/.env" ]]; then
  # shellcheck disable=SC1090
  source "${ROOT}/.env"
fi

if [[ -n "${DEEPSEEK_API_KEY:-}" ]]; then
  export BIOMNI_SOURCE="${BIOMNI_SOURCE:-Custom}"
  export BIOMNI_LLM="${BIOMNI_LLM:-${DEEPSEEK_MODEL_NAME:-deepseek-chat}}"
  export BIOMNI_CUSTOM_BASE_URL="${BIOMNI_CUSTOM_BASE_URL:-${DEEPSEEK_BASE_URL:-https://api.deepseek.com/v1}}"
  export BIOMNI_CUSTOM_API_KEY="${BIOMNI_CUSTOM_API_KEY:-${DEEPSEEK_API_KEY}}"
fi
if [[ -n "${BIOMNI_CUSTOM_BASE_URL:-}" && ! "${BIOMNI_CUSTOM_BASE_URL}" =~ ^https?:// ]]; then
  export BIOMNI_CUSTOM_BASE_URL="${DEEPSEEK_BASE_URL:-https://api.deepseek.com/v1}"
fi
if [[ -z "${ANTHROPIC_API_KEY:-}" && -z "${OPENAI_API_KEY:-}" && -z "${BIOMNI_CUSTOM_API_KEY:-}" ]]; then
  echo "Missing LLM API key. Set DEEPSEEK_API_KEY, BIOMNI_CUSTOM_API_KEY, OPENAI_API_KEY, or ANTHROPIC_API_KEY." >&2
  exit 2
fi
mkdir -p "${EXP_ROOT}/logs"

if [[ "${SHARD_COUNT}" -lt 1 ]]; then
  echo "--shard-count must be at least 1" >&2
  exit 2
fi

read -r -a SPLIT_ARRAY <<< "${SPLITS_TEXT}"

TMP_GRAPH_DIR=""
cleanup() {
  [[ -n "${TMP_GRAPH_DIR}" && -d "${TMP_GRAPH_DIR}" ]] && rm -rf "${TMP_GRAPH_DIR}"
}
trap cleanup EXIT

prepare_clean_graph() {
  local full_graph_dir="${BIOMANUS_LABBENCH_FULL_GRAPH_DIR:-}"
  if [[ -z "${full_graph_dir}" ]]; then
    full_graph_dir="$(find "${ROOT}/graph_outputs" -maxdepth 1 -type d -name 'mcp_generated_graph_all_*' | sort | tail -n 1)"
  fi
  if [[ -z "${full_graph_dir}" ]]; then
    "${PYTHON_BIN}" "${GRAPH_BUILDER}" \
      --preset all \
      --output-root "${ROOT}/graph_outputs" \
      --python-cmd "${PYTHON_BIN}" > /dev/null 2>&1
    full_graph_dir="$(find "${ROOT}/graph_outputs" -maxdepth 1 -type d -name 'mcp_generated_graph_all_*' | sort | tail -n 1)"
  fi
  if [[ -z "${full_graph_dir}" || ! -d "${full_graph_dir}" ]]; then
    echo "Could not locate or build a full MCP graph." >&2
    exit 2
  fi

  TMP_GRAPH_DIR="$(mktemp -d /tmp/biomanus_ablation_labbench_clean_graph_XXXXXX)"
  "${PYTHON_BIN}" "${GRAPH_SANITIZER}" \
    --source-graph-dir "${full_graph_dir}" \
    --out-dir "${TMP_GRAPH_DIR}" > "${EXP_ROOT}/logs/last_labbench_graph_sanitize.json"
  echo "Using clean LAB-Bench MCP graph: ${TMP_GRAPH_DIR}"
  cat "${TMP_GRAPH_DIR}/labbench_sanitize_summary.json"
}

prepare_clean_graph

run_variant() {
  local key="$1"
  local label="$2"
  shift 2
  local out_root="${EXP_ROOT}/results/${key}/labbench"
  local agent_root="${EXP_ROOT}/agent_runtime/${key}"
  local shard_log_root="${EXP_ROOT}/results/${key}/labbench_shard_logs"
  mkdir -p "${out_root}" "${agent_root}" "${shard_log_root}"

  echo "=== LAB-Bench ${label} ==="
  for eval_name in DbQA SeqQA; do
    local compact="${EXP_ROOT}/results/${key}/labbench_${eval_name}.jsonl"
    local reasoning="${EXP_ROOT}/results/${key}/labbench_${eval_name}_reasoning.log"
    touch "${compact}" "${reasoning}"

    local -a pids=()
    local -a labels=()
    for (( shard_index=0; shard_index<SHARD_COUNT; shard_index++ )); do
      local run_label="${key}_${eval_name}_shard$(printf '%02d' "$((shard_index + 1))")of$(printf '%02d' "${SHARD_COUNT}")"
      local shard_log="${shard_log_root}/${run_label}.log"
      local args=(
        --evals "${eval_name}"
        --splits "${SPLIT_ARRAY[@]}"
        --dev-size "${DEV_SIZE}"
        --test-size "${TEST_SIZE}"
        --output-root "${out_root}"
        --agent-root "${agent_root}"
        --graph-dir "${TMP_GRAPH_DIR}"
        --run-label "${run_label}"
        --shard-index "${shard_index}"
        --shard-count "${SHARD_COUNT}"
        --compact-results-path "${compact}"
        --reasoning-log-path "${reasoning}"
        --compact-output-only
        --skip-existing-results
        --timeout-seconds "${BIOMANUS_LABBENCH_TIMEOUT_SECONDS:-600}"
        "$@"
      )
      if [[ "${DEBUG}" -eq 1 ]]; then args+=(--debug); fi
      echo "[launcher] launching ${run_label} log=${shard_log}" | tee -a "${reasoning}"
      "${PYTHON_BIN}" "${RUNNER}" "${args[@]}" > "${shard_log}" 2>&1 &
      pids+=("$!")
      labels+=("${run_label}")
    done

    local failed=0
    for idx in "${!pids[@]}"; do
      local pid="${pids[$idx]}"
      local run_label="${labels[$idx]}"
      if wait "${pid}"; then
        echo "[launcher] shard ${run_label} completed" | tee -a "${reasoning}"
      else
        local status=$?
        failed=1
        echo "[launcher] shard ${run_label} failed exit_code=${status} log=${shard_log_root}/${run_label}.log" | tee -a "${reasoning}"
      fi
    done
    if [[ "${failed}" -ne 0 ]]; then
      return 1
    fi
  done
}

case "${VARIANT}" in
  biomanus)
    run_variant "biomanus" "full" --use-graph-retriever --use-mcp --mcp-retrieval-mode graph
    ;;
  mcp_flat)
    run_variant "mcp_flat" "MCP + flat retrieval" --use-graph-retriever --use-mcp --mcp-retrieval-mode flat
    ;;
  mcp_metadata)
    run_variant "mcp_metadata" "MCP + metadata retrieval" --use-graph-retriever --use-mcp --mcp-retrieval-mode metadata
    ;;
  minus_graph)
    run_variant "minus_graph" "minus graph" --no-graph-retriever --use-mcp
    ;;
  minus_mcp)
    run_variant "minus_mcp" "minus MCP" --use-graph-retriever --no-mcp
    ;;
  minus_mcp_graph)
    run_variant "minus_mcp_graph" "minus MCP and graph" --no-graph-retriever --no-mcp
    ;;
  all)
    run_variant "biomanus" "full" --use-graph-retriever --use-mcp --mcp-retrieval-mode graph
    run_variant "mcp_flat" "MCP + flat retrieval" --use-graph-retriever --use-mcp --mcp-retrieval-mode flat
    run_variant "mcp_metadata" "MCP + metadata retrieval" --use-graph-retriever --use-mcp --mcp-retrieval-mode metadata
    run_variant "minus_graph" "minus graph" --no-graph-retriever --use-mcp
    run_variant "minus_mcp" "minus MCP" --use-graph-retriever --no-mcp
    run_variant "minus_mcp_graph" "minus MCP and graph" --no-graph-retriever --no-mcp
    ;;
  *)
    echo "Unknown variant: ${VARIANT}" >&2
    exit 2
    ;;
esac