#!/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
mkdir -p outputs runs logs

TRAIN_JSONL="${TRAIN_JSONL:-data/processed/training_probe/train.jsonl}"
VAL_JSONL="${VAL_JSONL:-data/processed/training_probe/validation.jsonl}"
MAX_LENGTH="${MAX_LENGTH:-262144}"
SAVE_STEPS="${SAVE_STEPS:-1000}"
EVAL_STEPS="${EVAL_STEPS:-1000}"
LOGGING_STEPS="${LOGGING_STEPS:-1}"
NUM_EPOCHS="${NUM_EPOCHS:-1}"
NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
DEEPSPEED="${DEEPSPEED:-zero2}"
WARMUP_RATIO="${WARMUP_RATIO:-0.1}"
LR_SCHEDULER_TYPE="${LR_SCHEDULER_TYPE:-cosine}"
LORA_RANK="${LORA_RANK:-32}"
DEFAULT_PER_DEVICE_BATCH_SIZE="${DEFAULT_PER_DEVICE_BATCH_SIZE:-1}"
DEFAULT_GRAD_ACCUM_STEPS="${DEFAULT_GRAD_ACCUM_STEPS:-1}"
DEFAULT_EVAL_BATCH_SIZE="${DEFAULT_EVAL_BATCH_SIZE:-1}"

env_key() {
  printf '%s' "$1" | tr '[:lower:]-' '[:upper:]_' | sed 's/[^A-Z0-9_]/_/g'
}

env_or_default() {
  local name="$1"
  local fallback="$2"
  if [[ -n "${!name:-}" ]]; then
    printf '%s' "${!name}"
  else
    printf '%s' "${fallback}"
  fi
}

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

run_swift_train() {
  local model_path="$1"
  local train_type="$2"
  local run_name="$3"
  local output_dir="outputs/${run_name}"
  local tb_dir="runs/${run_name}"
  local log_file="logs/${run_name}.log"
  local run_key
  run_key="$(env_key "${run_name}")"

  local default_lr
  if [[ "${train_type}" == "lora" ]]; then
    default_lr="${LORA_LEARNING_RATE:-5e-5}"
  else
    default_lr="${FULL_LEARNING_RATE:-1e-5}"
  fi
  local learning_rate
  learning_rate="$(env_or_default "${run_key}_LEARNING_RATE" "${LEARNING_RATE:-${default_lr}}")"

  local type_key
  type_key="$(env_key "${train_type}")"
  local type_bsz_var="${type_key}_PER_DEVICE_BATCH_SIZE"
  local type_accum_var="${type_key}_GRAD_ACCUM_STEPS"
  local per_device_batch_size
  local grad_accum_steps
  local eval_batch_size
  per_device_batch_size="$(env_or_default "${run_key}_PER_DEVICE_BATCH_SIZE" "${PER_DEVICE_BATCH_SIZE:-${!type_bsz_var:-${DEFAULT_PER_DEVICE_BATCH_SIZE}}}")"
  grad_accum_steps="$(env_or_default "${run_key}_GRAD_ACCUM_STEPS" "${GRAD_ACCUM_STEPS:-${!type_accum_var:-${DEFAULT_GRAD_ACCUM_STEPS}}}")"
  eval_batch_size="$(env_or_default "${run_key}_EVAL_BATCH_SIZE" "${EVAL_PER_DEVICE_BATCH_SIZE:-${DEFAULT_EVAL_BATCH_SIZE}}")"

  require_file "${TRAIN_JSONL}"
  require_file "${VAL_JSONL}"
  mkdir -p "${output_dir}" "${tb_dir}" logs

  local cmd=(
    swift sft
    --model "${model_path}"
    --dataset "${TRAIN_JSONL}"
    --val_dataset "${VAL_JSONL}"
    --tuner_type "${train_type}"
    --torch_dtype bfloat16
    --num_train_epochs "${NUM_EPOCHS}"
    --per_device_train_batch_size "${per_device_batch_size}"
    --per_device_eval_batch_size "${eval_batch_size}"
    --gradient_accumulation_steps "${grad_accum_steps}"
    --learning_rate "${learning_rate}"
    --warmup_ratio "${WARMUP_RATIO}"
    --lr_scheduler_type "${LR_SCHEDULER_TYPE}"
    --max_length "${MAX_LENGTH}"
    --save_steps "${SAVE_STEPS}"
    --eval_steps "${EVAL_STEPS}"
    --logging_steps "${LOGGING_STEPS}"
    --report_to tensorboard
    --logging_dir "${tb_dir}"
    --output_dir "${output_dir}"
    --save_total_limit "${SAVE_TOTAL_LIMIT:-3}"
    --dataloader_num_workers "${DATALOADER_NUM_WORKERS:-4}"
    --deepspeed "${DEEPSPEED}"
  )

  if [[ "${train_type}" == "lora" ]]; then
    cmd+=(--lora_rank "${LORA_RANK}")
  fi
  if [[ -n "${MAX_STEPS:-}" ]]; then
    cmd+=(--max_steps "${MAX_STEPS}")
  fi

  printf 'CUDA_VISIBLE_DEVICES=%q NPROC_PER_NODE=%q ' "${CUDA_VISIBLE_DEVICES}" "${NPROC_PER_NODE}" | 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}" "${cmd[@]}" 2>&1 | tee "${log_file}"
}