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

# Visualize tracking JSON produced by exported video-root inference outputs.
#
# For each case, the script reads:
#   {case_dir}/merge.json
# and renders tracking boxes onto:
#   {VIDEO_ROOT_DIR}/{relative_case_dir}/sigmastar.1/camera4.bin
# then writes:
#   {case_dir}/tracking_vis/
#
# Supported modes:
#   1) Batch mode (default): walk RESULTS_ROOT and process all case dirs
#      containing TRACKING_JSON_NAME.
#   2) Root mode: pass an exported-infer results root as the first argument.
#   3) Single-case mode: pass a case dir or a direct tracking JSON path.
#
# Notes:
#   - The script uses --align-by-frame-id, so output is an image directory.
#   - TRACKING_JSON_NAME defaults to merge.json because combined_tracking.json
#     offsets track IDs by source and is less convenient for per-target review.

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}"

RESULTS_ROOT="${RESULTS_ROOT:-/data1/dongying/Mono3d/G1Q3/dataset_for_evaluation/CNCAP/CSTA_LN_outputs/model_20260416}"
VIDEO_ROOT_DIR="${VIDEO_ROOT_DIR:-/data1/dongying/Mono3d/G1Q3/dataset_for_evaluation/CNCAP/CSTA_LN}"
OUTPUT_ROOT="${OUTPUT_ROOT:-}"
OUTPUT_DIR_NAME="${OUTPUT_DIR_NAME:-tracking_vis}"
PYTHON_BIN="${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}"
VISUALIZE_SCRIPT="${VISUALIZE_SCRIPT:-${PROJECT_ROOT}/tools/temporal_analysis/visualize_tracking_boxes.py}"
TRACKING_JSON_NAME="${TRACKING_JSON_NAME:-merge.json}"
VIS_MAX_FRAMES="${VIS_MAX_FRAMES:-}"
VIS_CLASS_ID="${VIS_CLASS_ID:-}"
VIS_TRACK_IDS="${VIS_TRACK_IDS:-}"
VIS_SHOW_TRAJECTORY="${VIS_SHOW_TRAJECTORY:-0}"

TARGET_PATH=""
CLI_CLASS_ID=""
CLI_MAX_FRAMES=""
CLI_SHOW_TRAJECTORY=0
CLI_TRACK_IDS=()

while (($# > 0)); do
  case "$1" in
    --track-id)
      if (($# < 2)); then
        echo "Error: --track-id requires a value" >&2
        exit 1
      fi
      CLI_TRACK_IDS+=("$2")
      shift 2
      ;;
    --track-id=*)
      CLI_TRACK_IDS+=("${1#*=}")
      shift
      ;;
    --class-id)
      if (($# < 2)); then
        echo "Error: --class-id requires a value" >&2
        exit 1
      fi
      CLI_CLASS_ID="$2"
      shift 2
      ;;
    --class-id=*)
      CLI_CLASS_ID="${1#*=}"
      shift
      ;;
    --max-frames)
      if (($# < 2)); then
        echo "Error: --max-frames requires a value" >&2
        exit 1
      fi
      CLI_MAX_FRAMES="$2"
      shift 2
      ;;
    --max-frames=*)
      CLI_MAX_FRAMES="${1#*=}"
      shift
      ;;
    --show-trajectory)
      CLI_SHOW_TRAJECTORY=1
      shift
      ;;
    -*)
      echo "Error: unsupported option: $1" >&2
      exit 1
      ;;
    *)
      if [[ -n "${TARGET_PATH}" ]]; then
        echo "Error: multiple target paths provided: ${TARGET_PATH} and $1" >&2
        exit 1
      fi
      TARGET_PATH="$1"
      shift
      ;;
  esac
done

if [[ -z "${TARGET_PATH}" ]]; then
  TARGET_PATH="${RESULTS_ROOT}"
fi

if [[ ! -e "${TARGET_PATH}" ]]; then
  echo "Error: target path does not exist: ${TARGET_PATH}" >&2
  exit 1
fi

if [[ ! -d "${VIDEO_ROOT_DIR}" ]]; then
  echo "Error: VIDEO_ROOT_DIR does not exist: ${VIDEO_ROOT_DIR}" >&2
  exit 1
fi

resolve_abs_path() {
  local target_path="$1"
  if [[ -d "${target_path}" ]]; then
    (
      cd "${target_path}"
      pwd -P
    )
    return 0
  fi

  local parent_dir
  parent_dir=$(
    cd "$(dirname "${target_path}")"
    pwd -P
  )
  printf '%s/%s\n' "${parent_dir}" "$(basename "${target_path}")"
}

RESULTS_ROOT_ABS="$(resolve_abs_path "${RESULTS_ROOT}")"
VIDEO_ROOT_DIR_ABS="$(resolve_abs_path "${VIDEO_ROOT_DIR}")"

is_case_dir() {
  local dir_path="$1"
  [[ -d "${dir_path}" ]] && [[ -f "${dir_path}/${TRACKING_JSON_NAME}" ]]
}

resolve_case_dir() {
  local target_path="$1"
  if is_case_dir "${target_path}"; then
    printf '%s\n' "${target_path}"
    return 0
  fi
  if [[ -f "${target_path}" ]] && [[ "$(basename "${target_path}")" == "${TRACKING_JSON_NAME}" ]]; then
    dirname "${target_path}"
    return 0
  fi
  return 1
}

derive_output_dir() {
  local case_dir="$1"
  local case_abs
  local rel_case_dir

  case_abs="$(resolve_abs_path "${case_dir}")"
  if [[ -n "${OUTPUT_ROOT}" ]]; then
    if [[ "${case_abs}" == "${RESULTS_ROOT_ABS}" ]]; then
      printf '%s\n' "${OUTPUT_ROOT}"
      return 0
    fi
    if [[ "${case_abs}" == "${RESULTS_ROOT_ABS}"/* ]]; then
      rel_case_dir="${case_abs#"${RESULTS_ROOT_ABS}/"}"
      printf '%s/%s\n' "${OUTPUT_ROOT%/}" "${rel_case_dir}"
      return 0
    fi
    printf '%s/%s\n' "${OUTPUT_ROOT%/}" "$(basename "${case_abs}")"
    return 0
  fi

  printf '%s/%s\n' "${case_abs}" "${OUTPUT_DIR_NAME}"
}

derive_raw_video_path() {
  local case_dir="$1"
  local case_abs
  local rel_case_dir

  case_abs="$(resolve_abs_path "${case_dir}")"
  if [[ "${case_abs}" != "${RESULTS_ROOT_ABS}" && "${case_abs}" != "${RESULTS_ROOT_ABS}"/* ]]; then
    echo "Error: case directory is not under RESULTS_ROOT: ${case_abs}" >&2
    return 1
  fi

  if [[ "${case_abs}" == "${RESULTS_ROOT_ABS}" ]]; then
    echo "Error: cannot derive a raw video path from the results root itself" >&2
    return 1
  fi

  rel_case_dir="${case_abs#"${RESULTS_ROOT_ABS}/"}"
  printf '%s/%s/sigmastar.1/camera4.bin\n' "${VIDEO_ROOT_DIR_ABS}" "${rel_case_dir}"
}

build_track_id_args() {
  local -n out_ref=$1
  out_ref=()

  if [[ "${#CLI_TRACK_IDS[@]}" -gt 0 ]]; then
    for track_id in "${CLI_TRACK_IDS[@]}"; do
      out_ref+=(--track-ids "${track_id}")
    done
    return 0
  fi

  if [[ -n "${VIS_TRACK_IDS}" ]]; then
    # shellcheck disable=SC2206
    local track_id_arr=(${VIS_TRACK_IDS})
    for track_id in "${track_id_arr[@]}"; do
      out_ref+=(--track-ids "${track_id}")
    done
  fi
}

run_single_case() {
  local case_dir="$1"
  local tracking_json="${case_dir}/${TRACKING_JSON_NAME}"
  local raw_video_path
  local output_dir
  local cmd
  local track_id_args=()
  local class_id_value=""
  local max_frames_value=""
  local show_trajectory_value=0

  if [[ ! -f "${tracking_json}" ]]; then
    echo "[ERROR] ${TRACKING_JSON_NAME} not found in case directory: ${case_dir}" >&2
    return 1
  fi

  raw_video_path="$(derive_raw_video_path "${case_dir}")"
  if [[ ! -f "${raw_video_path}" ]]; then
    echo "[ERROR] camera4.bin not found for case: ${case_dir}" >&2
    echo "        expected: ${raw_video_path}" >&2
    return 1
  fi

  output_dir="$(derive_output_dir "${case_dir}")"
  build_track_id_args track_id_args

  if [[ -n "${CLI_CLASS_ID}" ]]; then
    class_id_value="${CLI_CLASS_ID}"
  elif [[ -n "${VIS_CLASS_ID}" ]]; then
    class_id_value="${VIS_CLASS_ID}"
  fi

  if [[ -n "${CLI_MAX_FRAMES}" ]]; then
    max_frames_value="${CLI_MAX_FRAMES}"
  elif [[ -n "${VIS_MAX_FRAMES}" ]]; then
    max_frames_value="${VIS_MAX_FRAMES}"
  fi

  if [[ "${CLI_SHOW_TRAJECTORY}" == "1" || "${VIS_SHOW_TRAJECTORY}" == "1" ]]; then
    show_trajectory_value=1
  fi

  cmd=(
    "${PYTHON_BIN}" "${VISUALIZE_SCRIPT}"
    --tracking "${tracking_json}"
    --images "${raw_video_path}"
    --output "${output_dir}"
    --align-by-frame-id
  )

  if [[ -n "${max_frames_value}" ]]; then
    cmd+=(--max-frames "${max_frames_value}")
  fi

  if [[ -n "${class_id_value}" ]]; then
    cmd+=(--class-id "${class_id_value}")
  fi

  if [[ "${#track_id_args[@]}" -gt 0 ]]; then
    cmd+=("${track_id_args[@]}")
  fi

  if [[ "${show_trajectory_value}" == "1" ]]; then
    cmd+=(--show-trajectory)
  fi

  echo ""
  echo "######################################################################"
  echo "# Tracking visualization for exported video-root outputs"
  echo "######################################################################"
  echo "Case   : ${case_dir}"
  echo "Track  : ${tracking_json}"
  echo "Video  : ${raw_video_path}"
  echo "Output : ${output_dir}"

  "${cmd[@]}"
}

run_batch_root() {
  local batch_root="$1"
  local total_cases=0
  local success_cases=0
  local failed_cases=0

  while IFS= read -r -d '' tracking_json; do
    local case_dir
    case_dir="$(dirname "${tracking_json}")"
    ((total_cases += 1))

    if run_single_case "${case_dir}"; then
      ((success_cases += 1))
    else
      ((failed_cases += 1))
      printf '[FAIL] case=%s\n' "${case_dir}" >&2
    fi
  done < <(
    find "${batch_root}" -type f -name "${TRACKING_JSON_NAME}" -print0 | sort -z
  )

  if [[ "${total_cases}" -eq 0 ]]; then
    echo "[ERROR] No ${TRACKING_JSON_NAME} files were found under: ${batch_root}" >&2
    return 1
  fi

  echo ""
  printf '[DONE] cases=%d success=%d failed=%d\n' \
    "${total_cases}" "${success_cases}" "${failed_cases}"

  [[ "${failed_cases}" -eq 0 ]]
}

if RESOLVED_CASE_DIR="$(resolve_case_dir "${TARGET_PATH}")"; then
  run_single_case "${RESOLVED_CASE_DIR}"
elif [[ -d "${TARGET_PATH}" ]]; then
  run_batch_root "${TARGET_PATH}"
else
  echo "Error: unsupported target path: ${TARGET_PATH}" >&2
  exit 1
fi
