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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
UMM_ROOT="${UMM_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}"
cd "${UMM_ROOT}"

unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY

export PYTHONPATH="${UMM_ROOT}/code/umm/runtime:${UMM_ROOT}/code/umm:${UMM_ROOT}/code/umm/stage2_covt:${UMM_ROOT}/third_party/transformers/src:${UMM_ROOT}/third_party/diffusers/src:${PYTHONPATH:-}"
export HF_HOME="${HF_HOME:-${UMM_ROOT}/cache/huggingface}"
export TORCH_HOME="${TORCH_HOME:-${UMM_ROOT}/cache/torch}"
export TMPDIR="${TMPDIR:-${UMM_ROOT}/tmp}"
export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
export TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-4}"

PYTHON_BIN="${PYTHON_BIN:-${UMM_ROOT}/.env/unvideo/bin/python}"
BENCHMARK="${BENCHMARK:?Set BENCHMARK to cvbench, blink, or vstar}"
CHECKPOINT_STEP="${CHECKPOINT_STEP:-7000}"
MODEL_PATH="${MODEL_PATH:-${UMM_ROOT}/models/ckpts/stage2_covt_from3k_bcd_2k/merged/checkpoint-${CHECKPOINT_STEP}-merged}"
MANIFEST="${MANIFEST:?Set MANIFEST to a prepared benchmark manifest}"
OUTPUT_DIR="${OUTPUT_DIR:?Set OUTPUT_DIR to a new result directory}"
MASTER_PORT="${MASTER_PORT:-29700}"
PROCESS_NAME="${PROCESS_NAME:-H3_${BENCHMARK}_7K}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-256}"
# SWANLAB_PROJECT is a reserved structured setting in SwanLab 0.9.x.  Accept
# it for backwards compatibility, then remove it from the Python environment.
SWANLAB_PROJECT_NAME="${SWANLAB_PROJECT_NAME:-${SWANLAB_PROJECT:-unified-model}}"
unset SWANLAB_PROJECT
SWANLAB_RUN_NAME="${SWANLAB_RUN_NAME:-H3_LoRA-ckpt${CHECKPOINT_STEP}-${BENCHMARK}}"
SWANLAB_KEY_FILE="${SWANLAB_KEY_FILE:-${UMM_ROOT}/secrets/swanlab_key.txt}"

if [[ ! -x "${PYTHON_BIN}" ]]; then
  echo "Missing Python environment: ${PYTHON_BIN}" >&2
  exit 1
fi
if [[ ! -f "${MODEL_PATH}/config.json" ]]; then
  echo "Merged model is not ready: ${MODEL_PATH}" >&2
  exit 1
fi
if [[ ! -s "${MANIFEST}" ]]; then
  echo "Manifest is missing: ${MANIFEST}" >&2
  exit 1
fi
if [[ -z "${SWANLAB_API_KEY:-}" ]]; then
  if [[ ! -s "${SWANLAB_KEY_FILE}" ]]; then
    echo "SwanLab key file is missing: ${SWANLAB_KEY_FILE}" >&2
    exit 1
  fi
  IFS= read -r SWANLAB_API_KEY < "${SWANLAB_KEY_FILE}" || true
  if [[ -z "${SWANLAB_API_KEY}" ]]; then
    echo "SwanLab key file is empty: ${SWANLAB_KEY_FILE}" >&2
    exit 1
  fi
  export SWANLAB_API_KEY
fi

exec "${PYTHON_BIN}" -m torch.distributed.run \
  --standalone \
  --nnodes=1 \
  --nproc-per-node=8 \
  "--master-port=${MASTER_PORT}" \
  "${UMM_ROOT}/code/umm/stage2_covt/eval_multibench.py" \
  --model-path "${MODEL_PATH}" \
  --manifest "${MANIFEST}" \
  --output-dir "${OUTPUT_DIR}" \
  --benchmark "${BENCHMARK}" \
  --checkpoint-step "${CHECKPOINT_STEP}" \
  --max-new-tokens "${MAX_NEW_TOKENS}" \
  --process-name "${PROCESS_NAME}" \
  --swanlab \
  --swanlab-project "${SWANLAB_PROJECT_NAME}" \
  --swanlab-run-name "${SWANLAB_RUN_NAME}" \
  "$@"