#!/bin/bash
set -Eeuo pipefail

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

RUN_ID="${RUN_ID:-$(date +%Y%m%d_%H%M%S)}"
PYTHON_BIN="${PYTHON_BIN:-/dev/shm/suiqk/conda_envs/scalerag-ts-v4/bin/python}"
ACCELERATE_CONFIG="${ACCELERATE_CONFIG:-/mnt/share01/sqk/ITFormer/accelerate_config.yaml}"
CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
MODEL_PATH="${MODEL_PATH:-/mnt/share01/sqk/models/ChatTime-1-7B-Chat}"
ALLOW_HF_DOWNLOAD="${ALLOW_HF_DOWNLOAD:-0}"
MQA_DIR="${MQA_DIR:-/mnt/share01/sqk/MQA}"
DATA_ROOT="${DATA_ROOT:-/mnt/share01/sqk/datasets/Time-MQA_TSQA/tmp}"
TRAIN_FILE="${TRAIN_FILE:-train.jsonl}"
EVAL_FILE="${EVAL_FILE:-eval.jsonl}"
OUTPUT_ROOT="${OUTPUT_ROOT:-${PROJECT_DIR}/tsqa_adapter/outputs/sft_${RUN_ID}}"
ADAPTER_OUTPUT_DIR="${ADAPTER_OUTPUT_DIR:-${OUTPUT_ROOT}/adapter}"
EVAL_OUTPUT_DIR="${EVAL_OUTPUT_DIR:-${OUTPUT_ROOT}/eval}"
RUN_EVAL_AFTER_SFT="${RUN_EVAL_AFTER_SFT:-1}"
MAX_TRAIN_SAMPLES="${MAX_TRAIN_SAMPLES:-}"
MAX_EVAL_SAMPLES="${MAX_EVAL_SAMPLES:-}"
EVAL_BATCH_SIZE="${EVAL_BATCH_SIZE:-4}"
MAX_SEQ_LENGTH="${MAX_SEQ_LENGTH:-4096}"
SKIP_OVERLONG="${SKIP_OVERLONG:-1}"
PER_DEVICE_TRAIN_BATCH_SIZE="${PER_DEVICE_TRAIN_BATCH_SIZE:-1}"
GRADIENT_ACCUMULATION_STEPS="${GRADIENT_ACCUMULATION_STEPS:-16}"
NUM_TRAIN_EPOCHS="${NUM_TRAIN_EPOCHS:-2}"
LEARNING_RATE="${LEARNING_RATE:-2e-4}"
LORA_RANK="${LORA_RANK:-16}"
LORA_ALPHA="${LORA_ALPHA:-32}"
LORA_DROPOUT="${LORA_DROPOUT:-0.05}"
SAVE_STEPS="${SAVE_STEPS:-200}"
LOGGING_STEPS="${LOGGING_STEPS:-10}"
SAVE_TOTAL_LIMIT="${SAVE_TOTAL_LIMIT:-2}"
DATALOADER_NUM_WORKERS="${DATALOADER_NUM_WORKERS:-4}"
LOAD_IN_4BIT="${LOAD_IN_4BIT:-0}"
GRADIENT_CHECKPOINTING="${GRADIENT_CHECKPOINTING:-1}"
TORCH_DTYPE="${TORCH_DTYPE:-fp16}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-256}"
MAX_INPUT_TOKENS="${MAX_INPUT_TOKENS:-3840}"
LOG_DIR="${LOG_DIR:-${PROJECT_DIR}/tsqa_adapter/logs}"
LOG_FILE="${LOG_FILE:-${LOG_DIR}/sft_4gpu_${RUN_ID}.log}"

mkdir -p "$LOG_DIR" "$OUTPUT_ROOT"

fail() {
    echo "$*" >&2
    exit 1
}

exec > >(tee -a "$LOG_FILE") 2>&1

export CUDA_VISIBLE_DEVICES
export PYTHONPATH="${PROJECT_DIR}:${MQA_DIR}:${PYTHONPATH:-}"
export MQA_DIR
export TOKENIZERS_PARALLELISM=false
export PYTHONWARNINGS="ignore::FutureWarning:transformers.utils.hub"
export WANDB_MODE=offline

# Force a guaranteed-local, writable temp dir. The repo (cwd) sits on a
# shared/NFS mount; if no standard temp dir is usable Python's tempfile falls
# back to the cwd, and DataLoader workers then try to bind their AF_UNIX IPC
# sockets there -- which NFS rejects with "Operation not permitted", crashing on
# the first batch and littering the repo with tmpXXXX scratch dirs. /dev/shm is
# local RAM, world-writable and supports unix sockets.
TMPDIR="${TMPDIR:-/dev/shm/tsqa_sft_tmp_${RUN_ID}}"
mkdir -p "$TMPDIR" || fail "Cannot create local TMPDIR: $TMPDIR (set TMPDIR=/path to a writable local dir)."
export TMPDIR TMP="$TMPDIR" TEMP="$TMPDIR"
trap 'rm -rf "$TMPDIR" 2>/dev/null || true' EXIT

[ -x "$PYTHON_BIN" ] || fail "Python executable not found: $PYTHON_BIN. Set PYTHON_BIN=/path/to/bin/python."
[ -f "$ACCELERATE_CONFIG" ] || fail "Accelerate config not found: $ACCELERATE_CONFIG"
[ -f "${DATA_ROOT}/${TRAIN_FILE}" ] || fail "TSQA train file not found: ${DATA_ROOT}/${TRAIN_FILE}"
[ -f "${DATA_ROOT}/${EVAL_FILE}" ] || fail "TSQA eval file not found: ${DATA_ROOT}/${EVAL_FILE}"
[ -n "$MODEL_PATH" ] || fail "MODEL_PATH is required. Use a local ChatTime model path, or set ALLOW_HF_DOWNLOAD=1 with a HuggingFace model id."

if [ ! -d "$MODEL_PATH" ] && [ "$ALLOW_HF_DOWNLOAD" != "1" ]; then
    fail "MODEL_PATH is not a local directory: $MODEL_PATH. Set ALLOW_HF_DOWNLOAD=1 if you intentionally want HuggingFace downloads."
fi

"$PYTHON_BIN" -c "import accelerate; print('accelerate:', accelerate.__version__)" || \
    fail "The selected Python cannot import accelerate: $PYTHON_BIN"
ACCELERATE_CMD=("$PYTHON_BIN" -m accelerate.commands.accelerate_cli)

if [ "$LOAD_IN_4BIT" = "1" ]; then
    "$PYTHON_BIN" -c "import importlib.metadata as m; print('bitsandbytes:', m.version('bitsandbytes'))" || \
        fail "LOAD_IN_4BIT=1 requires bitsandbytes in $PYTHON_BIN."
fi

TRAIN_EXTRA_ARGS=()
if [ -n "$MAX_TRAIN_SAMPLES" ]; then
    TRAIN_EXTRA_ARGS+=(--max_train_samples "$MAX_TRAIN_SAMPLES")
fi
if [ "$ALLOW_HF_DOWNLOAD" = "1" ]; then
    TRAIN_EXTRA_ARGS+=(--allow_hf_download)
fi
if [ "$LOAD_IN_4BIT" = "1" ]; then
    TRAIN_EXTRA_ARGS+=(--load_in_4bit)
fi
if [ "$GRADIENT_CHECKPOINTING" = "1" ]; then
    TRAIN_EXTRA_ARGS+=(--gradient_checkpointing)
fi
if [ "$SKIP_OVERLONG" = "1" ]; then
    TRAIN_EXTRA_ARGS+=(--skip_overlong)
else
    TRAIN_EXTRA_ARGS+=(--no_skip_overlong)
fi

echo "Run id: $RUN_ID"
echo "Python: $PYTHON_BIN"
echo "Accelerate: ${ACCELERATE_CMD[*]}"
echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES"
echo "Model path: $MODEL_PATH"
echo "Data root: $DATA_ROOT (train=$TRAIN_FILE eval=$EVAL_FILE)"
echo "Adapter output dir: $ADAPTER_OUTPUT_DIR"
echo "Eval output dir: $EVAL_OUTPUT_DIR"
echo "Log file: $LOG_FILE"

"${ACCELERATE_CMD[@]}" launch --config_file "$ACCELERATE_CONFIG" \
    tsqa_adapter/finetune_tsqa_lora.py \
    --model_path "$MODEL_PATH" \
    --data_root "$DATA_ROOT" \
    --train_file "$TRAIN_FILE" \
    --output_dir "$ADAPTER_OUTPUT_DIR" \
    --max_seq_length "$MAX_SEQ_LENGTH" \
    --per_device_train_batch_size "$PER_DEVICE_TRAIN_BATCH_SIZE" \
    --gradient_accumulation_steps "$GRADIENT_ACCUMULATION_STEPS" \
    --num_train_epochs "$NUM_TRAIN_EPOCHS" \
    --learning_rate "$LEARNING_RATE" \
    --lora_rank "$LORA_RANK" \
    --lora_alpha "$LORA_ALPHA" \
    --lora_dropout "$LORA_DROPOUT" \
    --save_steps "$SAVE_STEPS" \
    --logging_steps "$LOGGING_STEPS" \
    --save_total_limit "$SAVE_TOTAL_LIMIT" \
    --dataloader_num_workers "$DATALOADER_NUM_WORKERS" \
    --torch_dtype "$TORCH_DTYPE" \
    "${TRAIN_EXTRA_ARGS[@]}"

if [ "$RUN_EVAL_AFTER_SFT" = "1" ]; then
    EVAL_EXTRA_ARGS=()
    if [ -n "$MAX_EVAL_SAMPLES" ]; then
        EVAL_EXTRA_ARGS+=(--max_eval_samples "$MAX_EVAL_SAMPLES")
    fi
    if [ "$ALLOW_HF_DOWNLOAD" = "1" ]; then
        EVAL_EXTRA_ARGS+=(--allow_hf_download)
    fi

    "${ACCELERATE_CMD[@]}" launch --config_file "$ACCELERATE_CONFIG" \
        tsqa_adapter/eval_tsqa.py \
        --model_path "$MODEL_PATH" \
        --adapter_path "$ADAPTER_OUTPUT_DIR" \
        --data_root "$DATA_ROOT" \
        --eval_file "$EVAL_FILE" \
        --output_dir "$EVAL_OUTPUT_DIR" \
        --eval_batch_size "$EVAL_BATCH_SIZE" \
        --max_new_tokens "$MAX_NEW_TOKENS" \
        --max_input_tokens "$MAX_INPUT_TOKENS" \
        --torch_dtype "$TORCH_DTYPE" \
        "${EVAL_EXTRA_ARGS[@]}"
fi