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

export PYTHONPATH="$(pwd):${PYTHONPATH:-}"

# Submit N SageMaker webdatasetize jobs, one per split json:
#   <split_s3_prefix>/<split_prefix>_<i>.json, i=0..N-1
#
# Usage:
#   bash anydata/sagemaker/run_sagemaker_webdatasetize_splits.sh \
#     <dataset> <username> [instance_type] [queue] [num_splits] [first_build_type] [src_split_json]
#
# Example:
#   bash anydata/sagemaker/run_sagemaker_webdatasetize_splits.sh \
#     DDAD fzc g6e-small none 8
#   bash anydata/sagemaker/run_sagemaker_webdatasetize_splits.sh \
#     DDAD fzc g6e-small none 4 none \
#     s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/frames/DDAD/split_all.json

dataset=${1:?dataset required, e.g. DDAD}
name_prefix=${2:?username required, e.g. fzc}
instance_type=${3:-g6e-small}
queue_name=${4:-none}
num_splits=${5:-4}
first_build_type=${6:-${FIRST_BUILD_TYPE:-full}}

s3_bucket_root=${S3_BUCKET_ROOT:-s3://tri-ml-sandbox-16011-us-west-2-datasets}
args_file="anydata/sagemaker/webdatasetize_config/${dataset}.yaml"

cfg_values="$(python3 - "${args_file}" <<'PY'
import sys
import yaml

path = sys.argv[1]
with open(path, "r") as f:
    cfg = yaml.safe_load(f) or {}
mode = str(cfg.get("mode", "frames")).strip().lower()
from_mode = cfg.get("from_mode")
if from_mode is None:
    from_mode = mode if mode in {"frames", "videos"} else "frames"
s3_download_data = cfg.get("s3_download_data")
if not s3_download_data:
    raise SystemExit("Missing required 's3_download_data' in dataset YAML")
print(str(from_mode).strip().lower())
print(str(s3_download_data).strip())
PY
)"

from_mode="$(printf '%s\n' "${cfg_values}" | sed -n '1p')"
cfg_s3_download_data="$(printf '%s\n' "${cfg_values}" | sed -n '2p')"

split_s3_prefix=${dataset}
split_prefix=sm_split
unified_folder=${UNIFIED_FOLDER:-cv_unified/${from_mode}}
split_upload_s3_prefix=${SPLIT_UPLOAD_S3_PREFIX:-${s3_bucket_root}/${unified_folder}/${split_s3_prefix}}
if [[ "${cfg_s3_download_data}" == s3://* ]]; then
  default_src_split_json="${cfg_s3_download_data}"
else
  default_src_split_json="${s3_bucket_root}/${unified_folder}/${cfg_s3_download_data}"
fi
src_split_json=${7:-${default_src_split_json}}
if [[ "${src_split_json}" == "none" ]]; then
  src_split_json="${default_src_split_json}"
fi

rest_build_type=${REST_BUILD_TYPE:-none}
submit_sleep_s=${SUBMIT_SLEEP_S:-2}

echo "Submitting ${num_splits} jobs for dataset=${dataset}"
echo "split source pattern: ${split_s3_prefix}/${split_prefix}_<i>.json"
echo "unified folder: ${unified_folder}"
echo "build types: first=${first_build_type}, rest=${rest_build_type}"

echo
echo "Generating split files from: ${src_split_json}"
echo "Uploading generated splits to: ${split_upload_s3_prefix}"
python3 anydata/sagemaker/split_split_file.py \
  "${src_split_json}" \
  --dataset "${dataset}" \
  --num_splits "${num_splits}" \
  --prefix "${split_prefix}" \
  --upload_s3_prefix "${split_upload_s3_prefix}"

for ((i=0; i<num_splits; i++)); do
  split_json="${split_s3_prefix}/${split_prefix}_${i}.json"
  run_name="${name_prefix}-${dataset}-${split_prefix}_${i}"

  if [[ ${i} -eq 0 ]]; then
    export BUILD_TYPE="${first_build_type}"
  else
    export BUILD_TYPE="${rest_build_type}"
  fi

  echo
  echo "[$(printf "%02d" "$((i + 1))")/${num_splits}] submit ${run_name} with BUILD_TYPE=${BUILD_TYPE}"
  echo "split_json=${split_json}"

  bash anydata/sagemaker/run_sagemaker_webdatasetize.sh \
    "${dataset}" \
    "${run_name}" \
    "${instance_type}" \
    "${queue_name}" \
    "${name_prefix}" \
    "${split_json}"

  sleep "${submit_sleep_s}"
done

echo
echo "All submissions requested."