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

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

# Parallel launcher for event-json inference by scene.
#
# Strategy:
#   1. Read scene names from the top-level keys of EVENT_JSON_FILE.
#   2. Distribute scenes across GPU_IDS using round-robin or contiguous mode.
#   3. Each scene runs the existing event-json launcher in an isolated temp root.
#   4. After a scene succeeds, move scene artifacts into the final roots.
#
# This avoids concurrent writes to the same root-level manifest/status files while
# preserving the existing single-scene inference, tracking, and conversion logic.

WORKER_SCRIPT="${SCRIPT_DIR}/run_two_roi_exported_onnx_infer_event_json.sh"
PYTHON_BIN="${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}"
EVENT_JSON_FILE="${EVENT_JSON_FILE:-${MODEL_INFERENCE_DIR}/examples/events/G1Q3_场地评测数据集.json}"
EVENT_ID_FIELD="${EVENT_ID_FIELD:-event_uuid}"
EVENT_CLIP_IDS_FIELD="${EVENT_CLIP_IDS_FIELD:-clips}"
GPU_IDS="${GPU_IDS:-0 1}"
SCENES="${SCENES:-}"
SCENE_REGEX="${SCENE_REGEX:-}"
MAX_SCENES="${MAX_SCENES:-0}"
SCENE_DISTRIBUTION="${SCENE_DISTRIBUTION:-round_robin}"
DRY_RUN="${DRY_RUN:-0}"
OUTPUT_DIR="${OUTPUT_DIR:-/data1/dongying/Mono3d/G1Q3/model_inference/cncap/cncap_2021_with_cls_ego/model_20260506_keep_fake3d_validate}"
EXPORT_ROOT="${EXPORT_ROOT:-/data1/dongying/Mono3d/G1Q3/model_inference/cncap/cncap_2021_with_cls_ego/event_exports_new}"
CONVERT_OUTPUT_ROOT_BASE="${CONVERT_OUTPUT_ROOT:-${OUTPUT_DIR%/}_objectlist}"
PARALLEL_RUN_ID="${PARALLEL_RUN_ID:-$(date +%Y%m%d_%H%M%S)}"
PARALLEL_TMP_ROOT="${PARALLEL_TMP_ROOT:-${OUTPUT_DIR%/}/_parallel_scene_tmp/${PARALLEL_RUN_ID}}"
PARALLEL_LOG_DIR="${PARALLEL_LOG_DIR:-${OUTPUT_DIR%/}/parallel_scene_logs/${PARALLEL_RUN_ID}}"
SKIP_EXISTING_SCENES="${SKIP_EXISTING_SCENES:-0}"
FORWARD_ARGS=("$@")

discover_scenes() {
  if [[ -n "${SCENES}" ]]; then
    # shellcheck disable=SC2206
    DISCOVERED_SCENES=(${SCENES})
  else
    mapfile -t DISCOVERED_SCENES < <(
      "${PYTHON_BIN}" - "${EVENT_JSON_FILE}" <<'PY'
import json
import sys

json_path = sys.argv[1]
with open(json_path, encoding="utf-8") as file:
    payload = json.load(file)
if not isinstance(payload, dict):
    raise SystemExit(f"Expected top-level object in {json_path}, got {type(payload).__name__}")
for scene_name in sorted(payload):
    print(scene_name)
PY
    )
  fi

  if [[ -n "${SCENE_REGEX}" ]]; then
    local -a filtered_scenes=()
    local scene_name
    for scene_name in "${DISCOVERED_SCENES[@]}"; do
      if [[ "${scene_name}" =~ ${SCENE_REGEX} ]]; then
        filtered_scenes+=("${scene_name}")
      fi
    done
    DISCOVERED_SCENES=("${filtered_scenes[@]}")
  fi

  if [[ "${MAX_SCENES}" != "0" ]] && [[ "${MAX_SCENES}" -lt "${#DISCOVERED_SCENES[@]}" ]]; then
    DISCOVERED_SCENES=("${DISCOVERED_SCENES[@]:0:${MAX_SCENES}}")
  fi

  if [[ "${#DISCOVERED_SCENES[@]}" -eq 0 ]]; then
    echo "No scenes selected from ${EVENT_JSON_FILE}" >&2
    exit 1
  fi
}

scene_assigned_to_worker() {
  local scene_index="$1"
  local worker_index="$2"
  local worker_count="$3"

  case "${SCENE_DISTRIBUTION}" in
    round_robin)
      (( scene_index % worker_count == worker_index ))
      ;;
    contiguous)
      local start=$(( ${#DISCOVERED_SCENES[@]} * worker_index / worker_count ))
      local end=$(( ${#DISCOVERED_SCENES[@]} * (worker_index + 1) / worker_count ))
      (( scene_index >= start && scene_index < end ))
      ;;
    *)
      echo "Unsupported SCENE_DISTRIBUTION: ${SCENE_DISTRIBUTION}" >&2
      exit 1
      ;;
  esac
}

print_scene_plan() {
  local worker_count="$1"
  local worker_index
  local scene_index
  local scene_name

  echo ""
  echo "######################################################################"
  echo "# Dry Run: scene to GPU assignment"
  echo "######################################################################"
  echo "Event JSON       : ${EVENT_JSON_FILE}"
  echo "Event id field   : ${EVENT_ID_FIELD}"
  echo "Clip ids field   : ${EVENT_CLIP_IDS_FIELD}"
  echo "Scenes           : ${#DISCOVERED_SCENES[@]}"
  echo "GPU ids          : ${GPU_IDS}"
  echo "Distribution     : ${SCENE_DISTRIBUTION}"
  echo "Output root      : ${OUTPUT_DIR}"
  echo "Export root      : ${EXPORT_ROOT}"
  echo "Convert root     : ${CONVERT_OUTPUT_ROOT_BASE}"
  echo "Temp root        : ${PARALLEL_TMP_ROOT}"
  echo "Worker logs      : ${PARALLEL_LOG_DIR}"

  for worker_index in "${!GPU_IDS_ARR[@]}"; do
    echo ""
    echo "[worker ${worker_index}] gpu=${GPU_IDS_ARR[${worker_index}]}"
    for scene_index in "${!DISCOVERED_SCENES[@]}"; do
      if ! scene_assigned_to_worker "${scene_index}" "${worker_index}" "${worker_count}"; then
        continue
      fi
      scene_name="${DISCOVERED_SCENES[${scene_index}]}"
      echo "  - ${scene_name}"
    done
  done
}

move_scene_tree() {
  local src_dir="$1"
  local dst_dir="$2"

  [[ -d "${src_dir}" ]] || return 0
  if [[ -e "${dst_dir}" ]]; then
    echo "Target already exists, refusing to overwrite: ${dst_dir}" >&2
    return 1
  fi

  mkdir -p "$(dirname "${dst_dir}")"
  mv "${src_dir}" "${dst_dir}"
}

run_one_scene() {
  local gpu_id="$1"
  local worker_index="$2"
  local scene_name="$3"
  local scene_tmp_root="${PARALLEL_TMP_ROOT%/}/worker_${worker_index}_gpu_${gpu_id}/${scene_name}"
  local worker_output_root="${scene_tmp_root}/output_root"
  local worker_export_root="${scene_tmp_root}/export_root"
  local worker_convert_root="${scene_tmp_root}/convert_root"
  local final_scene_output_dir="${OUTPUT_DIR%/}/${scene_name}"
  local final_scene_export_dir="${EXPORT_ROOT%/}/${scene_name}"
  local final_scene_convert_dir="${CONVERT_OUTPUT_ROOT_BASE%/}/${scene_name}"
  local output_status_dir="${OUTPUT_DIR%/}/_status"

  if [[ "${SKIP_EXISTING_SCENES}" == "1" ]] && [[ -d "${final_scene_output_dir}" ]]; then
    echo "[GPU ${gpu_id}] skip existing scene: ${scene_name}"
    return 0
  fi

  mkdir -p "${scene_tmp_root}"
  echo "[GPU ${gpu_id}] start scene: ${scene_name}"

  CUDA_VISIBLE_DEVICES="${gpu_id}" \
  SCENE="${scene_name}" \
  EVENT_ID_FIELD="${EVENT_ID_FIELD}" \
  EVENT_CLIP_IDS_FIELD="${EVENT_CLIP_IDS_FIELD}" \
  OUTPUT_DIR="${worker_output_root}" \
  EXPORT_ROOT="${worker_export_root}" \
  CONVERT_OUTPUT_ROOT="${worker_convert_root}" \
  bash "${WORKER_SCRIPT}" "${FORWARD_ARGS[@]}"

  move_scene_tree "${worker_output_root}/${scene_name}" "${final_scene_output_dir}"
  move_scene_tree "${worker_export_root}/${scene_name}" "${final_scene_export_dir}"

  if [[ "${CONVERT_OUTPUT_LAYOUT:-parallel_root}" == "parallel_root" ]]; then
    move_scene_tree "${worker_convert_root}/${scene_name}" "${final_scene_convert_dir}"
  fi

  if [[ -f "${worker_output_root}/_status/event_scene_manifest.json" ]]; then
    mkdir -p "${output_status_dir}"
    cp "${worker_output_root}/_status/event_scene_manifest.json" \
      "${output_status_dir}/${scene_name}_event_scene_manifest.json"
  fi

  echo "[GPU ${gpu_id}] finished scene: ${scene_name}"
}

launch_worker() {
  local worker_index="$1"
  local gpu_id="$2"
  local worker_count="$3"
  local log_file="${PARALLEL_LOG_DIR%/}/worker_${worker_index}_gpu_${gpu_id}.log"
  local scene_index
  local scene_name
  local worker_failed=0
  local worker_scene_count=0

  for scene_index in "${!DISCOVERED_SCENES[@]}"; do
    if scene_assigned_to_worker "${scene_index}" "${worker_index}" "${worker_count}"; then
      ((worker_scene_count+=1))
    fi
  done

  echo "Launch worker ${worker_index} on GPU ${gpu_id}: ${worker_scene_count} scenes -> ${log_file}"
  (
    echo "Worker index     : ${worker_index}"
    echo "GPU id           : ${gpu_id}"
    echo "Scene count      : ${worker_scene_count}"
    echo "Distribution     : ${SCENE_DISTRIBUTION}"
    echo "Event JSON       : ${EVENT_JSON_FILE}"
    echo "Event id field   : ${EVENT_ID_FIELD}"
    echo "Clip ids field   : ${EVENT_CLIP_IDS_FIELD}"
    echo "Output root      : ${OUTPUT_DIR}"
    echo "Export root      : ${EXPORT_ROOT}"
    echo "Convert root     : ${CONVERT_OUTPUT_ROOT_BASE}"
    echo "Temp root        : ${PARALLEL_TMP_ROOT}"

    for scene_index in "${!DISCOVERED_SCENES[@]}"; do
      if ! scene_assigned_to_worker "${scene_index}" "${worker_index}" "${worker_count}"; then
        continue
      fi
      scene_name="${DISCOVERED_SCENES[${scene_index}]}"
      if ! run_one_scene "${gpu_id}" "${worker_index}" "${scene_name}"; then
        worker_failed=1
        echo "[GPU ${gpu_id}] failed scene: ${scene_name}" >&2
      fi
    done

    exit "${worker_failed}"
  ) >"${log_file}" 2>&1 &

  WORKER_PIDS+=("$!")
  WORKER_LOGS+=("${log_file}")
}

if [[ ! -f "${WORKER_SCRIPT}" ]]; then
  echo "Worker script not found: ${WORKER_SCRIPT}" >&2
  exit 1
fi

discover_scenes

# shellcheck disable=SC2206
GPU_IDS_ARR=(${GPU_IDS})
if [[ "${#GPU_IDS_ARR[@]}" -lt 1 ]]; then
  echo "GPU_IDS must contain at least one GPU id." >&2
  exit 1
fi

worker_count="${#GPU_IDS_ARR[@]}"

if [[ "${DRY_RUN}" == "1" ]]; then
  print_scene_plan "${worker_count}"
  exit 0
fi

mkdir -p "${PARALLEL_LOG_DIR}" "${OUTPUT_DIR}" "${EXPORT_ROOT}"
if [[ "${CONVERT_OUTPUT_LAYOUT:-parallel_root}" == "parallel_root" ]]; then
  mkdir -p "${CONVERT_OUTPUT_ROOT_BASE}"
fi

echo ""
echo "######################################################################"
echo "# Parallel event-json inference by scene"
echo "######################################################################"
echo "Event JSON       : ${EVENT_JSON_FILE}"
echo "Event id field   : ${EVENT_ID_FIELD}"
echo "Clip ids field   : ${EVENT_CLIP_IDS_FIELD}"
echo "Scenes           : ${#DISCOVERED_SCENES[@]}"
echo "GPU ids          : ${GPU_IDS}"
echo "Distribution     : ${SCENE_DISTRIBUTION}"
echo "Output root      : ${OUTPUT_DIR}"
echo "Export root      : ${EXPORT_ROOT}"
echo "Convert root     : ${CONVERT_OUTPUT_ROOT_BASE}"
echo "Temp root        : ${PARALLEL_TMP_ROOT}"
echo "Worker logs      : ${PARALLEL_LOG_DIR}"

WORKER_PIDS=()
WORKER_LOGS=()
for worker_index in "${!GPU_IDS_ARR[@]}"; do
  launch_worker "${worker_index}" "${GPU_IDS_ARR[${worker_index}]}" "${worker_count}"
done

failed=0
for worker_index in "${!WORKER_PIDS[@]}"; do
  if wait "${WORKER_PIDS[${worker_index}]}"; then
    echo "Worker ${worker_index}/${worker_count} finished successfully. Log: ${WORKER_LOGS[${worker_index}]}"
  else
    echo "[ERROR] Worker ${worker_index}/${worker_count} failed. Log: ${WORKER_LOGS[${worker_index}]}" >&2
    failed=1
  fi
done

if [[ "${failed}" == "1" ]]; then
  exit 1
fi

echo "All scene workers finished successfully."