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

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

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": "video-finetuner",
  "videoModel": "${VIDEO_MODEL:-}",
  "datasetPath": "${DATASET_PATH:-}"
}
EOF
}

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

trap on_exit EXIT

mkdir -p "$artifacts_dir" "$output_subdir"
write_status "running" "Preparing video finetuning job"

if [ "${SIMULATE_ONLY:-1}" = "1" ]; then
  sleep "${SIMULATED_DURATION_SECONDS:-8}"
  cat >"${output_subdir}/training-plan.json" <<EOF
{
  "model": "${VIDEO_MODEL:-THUDM/CogVideoX-5b}",
  "dataset_path": "${DATASET_PATH:-/datasets/video-captions}",
  "max_train_steps": "${MAX_STEPS:-1000}",
  "suggested_stack": ["diffusers", "accelerate", "peft"]
}
EOF
  cat >"${artifacts_dir}/run-summary.txt" <<EOF
Simulation mode completed.

Replace TRAIN_COMMAND and set SIMULATE_ONLY=0 when your container contains the Diffusers / Accelerate video training stack.
Suggested starting point:
${TRAIN_COMMAND:-accelerate launch train_video_lora.py ...}
EOF
  write_status "completed" "Simulation complete; placeholder video finetuning manifest 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
Video finetuning command completed successfully.
Model: ${VIDEO_MODEL:-}
Dataset path: ${DATASET_PATH:-}
Output directory: ${output_subdir}
EOF

write_status "completed" "Video finetuning command completed successfully"