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

status_file="${PWD}/status.json"
artifacts_dir="${PWD}/artifacts"
output_subdir="${OUTPUT_SUBDIR:-artifacts/lora-adapter}"

timestamp() {
  date -u +"%Y-%m-%dT%H:%M:%SZ"
}

write_status() {
  local state="$1"
  local message="$2"
  cat >"$status_file" <<EOF
{
  "status": "$state",
  "message": "$message",
  "updatedAt": "$(timestamp)",
  "template": "lora-trainer",
  "baseModel": "${BASE_MODEL:-}",
  "dataset": "${DATASET_NAME:-}"
}
EOF
}

on_exit() {
  local rc=$?
  if [ "$rc" -ne 0 ]; then
    write_status "failed" "train.sh exited with status ${rc}"
  fi
}

trap on_exit EXIT

mkdir -p "$artifacts_dir" "$output_subdir"
write_status "running" "Preparing LoRA / QLoRA training run"

if [ "${SIMULATE_ONLY:-1}" = "1" ]; then
  sleep "${SIMULATED_DURATION_SECONDS:-5}"
  cat >"${output_subdir}/adapter_config.json" <<EOF
{
  "base_model_name_or_path": "${BASE_MODEL:-meta-llama/Llama-3.1-8B-Instruct}",
  "peft_type": "LORA",
  "r": 16,
  "lora_alpha": 32,
  "target_modules": ["q_proj", "k_proj", "v_proj", "o_proj"]
}
EOF
  cat >"${artifacts_dir}/run-summary.txt" <<EOF
Simulation mode completed.

Replace TRAIN_COMMAND and set SIMULATE_ONLY=0 when your container has PEFT / Unsloth installed.
Suggested starting point:
${TRAIN_COMMAND:-python3 -m unsloth.cli.train ...}
EOF
  write_status "completed" "Simulation complete; placeholder adapter metadata written to artifacts"
  exit 0
fi

: "${TRAIN_COMMAND:?TRAIN_COMMAND must be set when SIMULATE_ONLY=0}"
bash -lc "${TRAIN_COMMAND}"

cat >"${artifacts_dir}/run-summary.txt" <<EOF
Training command completed successfully.
Base model: ${BASE_MODEL:-}
Dataset: ${DATASET_NAME:-}
Output directory: ${output_subdir}
EOF

write_status "completed" "Training command completed successfully"