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

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT_DIR}"

export http_proxy="${http_proxy:-http://100.72.0.101:8888}"
export https_proxy="${https_proxy:-http://100.72.0.101:8888}"
export HTTP_PROXY="${HTTP_PROXY:-${http_proxy}}"
export HTTPS_PROXY="${HTTPS_PROXY:-${https_proxy}}"
export HF_ENDPOINT="${HF_ENDPOINT:-https://hf-mirror.com}"
export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}"
export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"

if [[ -f .venv/bin/activate ]]; then
  source .venv/bin/activate
elif [[ "${DRY_RUN:-0}" != "1" ]]; then
  echo "Missing .venv. Run ./scripts/setup_env.sh first." >&2
  exit 2
fi

BEEGFS_ROOT="${BEEGFS_ROOT:-/mnt/beegfs/workspace/ti_coding_agent_probe}"
export MODELSCOPE_CACHE="${MODELSCOPE_CACHE:-${BEEGFS_ROOT}/modelscope_cache}"
MEGATRON_OUTPUT_ROOT="${MEGATRON_OUTPUT_ROOT:-${BEEGFS_ROOT}/megatron_outputs}"
MEGATRON_LOG_ROOT="${MEGATRON_LOG_ROOT:-${BEEGFS_ROOT}/megatron_logs}"
mkdir -p "${MODELSCOPE_CACHE}" "${MEGATRON_OUTPUT_ROOT}" "${MEGATRON_LOG_ROOT}" logs

TRAIN_DATASET="${TRAIN_DATASET:-${TRAIN_JSONL:-data/processed/training_probe/train.jsonl}}"
VAL_DATASET="${VAL_DATASET:-${VAL_JSONL:-data/processed/training_probe/validation.jsonl}}"
MEGATRON_MODEL="${MEGATRON_MODEL:-${QWEN36_27B_MODEL_ID:-Qwen/Qwen3.6-27B}}"
MEGATRON_RUN_NAME="${MEGATRON_RUN_NAME:-qwen36_27b_megatron_full}"
MEGATRON_OUTPUT_DIR="${MEGATRON_OUTPUT_DIR:-${MEGATRON_OUTPUT_ROOT}/${MEGATRON_RUN_NAME}}"

NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
TENSOR_MODEL_PARALLEL_SIZE="${TENSOR_MODEL_PARALLEL_SIZE:-4}"
PIPELINE_MODEL_PARALLEL_SIZE="${PIPELINE_MODEL_PARALLEL_SIZE:-1}"
CONTEXT_PARALLEL_SIZE="${CONTEXT_PARALLEL_SIZE:-1}"
MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-1}"
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-8}"
MAX_LENGTH="${MAX_LENGTH:-262144}"
LR="${LR:-1e-5}"
MIN_LR="${MIN_LR:-1e-6}"
LR_WARMUP_FRACTION="${LR_WARMUP_FRACTION:-0.1}"
NUM_EPOCHS="${NUM_EPOCHS:-1}"
SAVE_STEPS="${SAVE_STEPS:-1000}"
EVAL_STEPS="${EVAL_STEPS:-1000}"
DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-4}"
DATASET_NUM_PROC="${DATASET_NUM_PROC:-4}"
RECOMPUTE_GRANULARITY="${RECOMPUTE_GRANULARITY:-full}"
RECOMPUTE_METHOD="${RECOMPUTE_METHOD:-uniform}"
RECOMPUTE_NUM_LAYERS="${RECOMPUTE_NUM_LAYERS:-1}"
CROSS_ENTROPY_LOSS_FUSION="${CROSS_ENTROPY_LOSS_FUSION:-true}"
GRADIENT_ACCUMULATION_FUSION="${GRADIENT_ACCUMULATION_FUSION:-false}"
SAVE_SAFETENSORS="${SAVE_SAFETENSORS:-true}"
SEQUENCE_PARALLEL="${SEQUENCE_PARALLEL:-true}"

require_file() {
  if [[ ! -f "$1" ]]; then
    echo "Missing required file: $1" >&2
    exit 2
  fi
}

run_megatron_sft() {
  require_file "${TRAIN_DATASET}"
  require_file "${VAL_DATASET}"
  mkdir -p "${MEGATRON_OUTPUT_DIR}"

  local log_file="${MEGATRON_LOG_ROOT}/${MEGATRON_RUN_NAME}.log"
  local cmd=(
    megatron sft
    --model "${MEGATRON_MODEL}"
    --save_safetensors "${SAVE_SAFETENSORS}"
    --dataset "${TRAIN_DATASET}"
    --val_dataset "${VAL_DATASET}"
    --tensor_model_parallel_size "${TENSOR_MODEL_PARALLEL_SIZE}"
    --pipeline_model_parallel_size "${PIPELINE_MODEL_PARALLEL_SIZE}"
    --context_parallel_size "${CONTEXT_PARALLEL_SIZE}"
    --sequence_parallel "${SEQUENCE_PARALLEL}"
    --micro_batch_size "${MICRO_BATCH_SIZE}"
    --global_batch_size "${GLOBAL_BATCH_SIZE}"
    --recompute_granularity "${RECOMPUTE_GRANULARITY}"
    --recompute_method "${RECOMPUTE_METHOD}"
    --recompute_num_layers "${RECOMPUTE_NUM_LAYERS}"
    --finetune true
    --cross_entropy_loss_fusion "${CROSS_ENTROPY_LOSS_FUSION}"
    --gradient_accumulation_fusion "${GRADIENT_ACCUMULATION_FUSION}"
    --lr "${LR}"
    --lr_warmup_fraction "${LR_WARMUP_FRACTION}"
    --min_lr "${MIN_LR}"
    --num_train_epochs "${NUM_EPOCHS}"
    --output_dir "${MEGATRON_OUTPUT_DIR}"
    --save_steps "${SAVE_STEPS}"
    --eval_steps "${EVAL_STEPS}"
    --max_length "${MAX_LENGTH}"
    --dataloader_num_workers "${DATALOADER_NUM_WORKERS}"
    --dataset_num_proc "${DATASET_NUM_PROC}"
    --no_save_optim true
    --no_save_rng true
    --model_author swift
    --model_name ti-coding-agent
  )

  printf 'CUDA_VISIBLE_DEVICES=%q NPROC_PER_NODE=%q MODELSCOPE_CACHE=%q ' \
    "${CUDA_VISIBLE_DEVICES}" "${NPROC_PER_NODE}" "${MODELSCOPE_CACHE}" | tee "${log_file}.cmd"
  printf '%q ' "${cmd[@]}" | tee -a "${log_file}.cmd"
  echo | tee -a "${log_file}.cmd"
  if [[ "${DRY_RUN:-0}" == "1" ]]; then
    return 0
  fi
  CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" \
    NPROC_PER_NODE="${NPROC_PER_NODE}" \
    MODELSCOPE_CACHE="${MODELSCOPE_CACHE}" \
    "${cmd[@]}" 2>&1 | tee "${log_file}"
}