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

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

DOWNLOAD_ROOT=${DOWNLOAD_ROOT:-/data1/dongying/Mono3d/G1Q3/feishu_project/downloaded_issue_data}
OUTPUT_ROOT=${OUTPUT_ROOT:-/data1/dongying/Mono3d/G1Q3/feishu_project/inference_issue_data}
MANIFEST_PATH=${MANIFEST_PATH:-"${OUTPUT_ROOT}/inference_manifest.json"}
PYTHON_BIN=${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}
INFERENCE_SCRIPT=${INFERENCE_SCRIPT:-"${PROJECT_ROOT}/tools/model_inference/core/run_two_roi_exported_onnx_infer.py"}
ISSUE_TRACKING_SCRIPT=${ISSUE_TRACKING_SCRIPT:-"${PROJECT_ROOT}/tools/feishu_project/run_issue_data_tracking.sh"}
ISSUE_JSON=${ISSUE_JSON:-}
VIDEO_STRIDE=${VIDEO_STRIDE:-1}
MAX_IMAGES=${MAX_IMAGES:-0}
FRAME_INDEX_START=${FRAME_INDEX_START:-}
FRAME_INDEX_END=${FRAME_INDEX_END:-}
FRAME_ID_START=${FRAME_ID_START:-}
FRAME_ID_END=${FRAME_ID_END:-}
TARGET_FRAME_ID=${TARGET_FRAME_ID:-}
FRAME_BEFORE=${FRAME_BEFORE:-100}
FRAME_AFTER=${FRAME_AFTER:-100}
USE_ISSUE_FRAME_WINDOW=${USE_ISSUE_FRAME_WINDOW:-0}
MISSING_ISSUE_FRAME_POLICY=${MISSING_ISSUE_FRAME_POLICY:-full}
EXPORTED_MODEL=${EXPORTED_MODEL:-${PROJECT_ROOT}/runs/export/train_mono3d_two_roi_20260416-raw_no_edge/merged_model.torchscript}
DEVICE=${DEVICE:-}
PROVIDERS=${PROVIDERS:-}
ENABLE_ATTR=${ENABLE_ATTR:-0}
ENABLE_CROSS_CLASS_MERGE_PRIOR=${ENABLE_CROSS_CLASS_MERGE_PRIOR:-1}
SAVE_AGGREGATE_PREDICTIONS=${SAVE_AGGREGATE_PREDICTIONS:-1}
SKIP_EXISTING=${SKIP_EXISTING:-0}
DRY_RUN=${DRY_RUN:-0}
ENABLE_TRACKING=${ENABLE_TRACKING:-0}
INFERENCE_EXTRA_ARGS=${INFERENCE_EXTRA_ARGS:-}

resolve_tracking_model_version() {
  if [[ -n "${TRACKING_MODEL_VERSION:-}" ]]; then
    printf '%s\n' "${TRACKING_MODEL_VERSION}"
    return 0
  fi

  if [[ -n "${EXPORTED_MODEL}" ]] && [[ "${EXPORTED_MODEL}" =~ ([0-9]{8}) ]]; then
    printf '%s\n' "${BASH_REMATCH[1]}"
    return 0
  fi

  return 1
}

FORWARD_ARGS=()
ISSUE_ID_ARGS=()
while (($# > 0)); do
  case "$1" in
    --issue-id)
      if (($# < 2)); then
        echo "Error: --issue-id requires a value" >&2
        exit 1
      fi
      ISSUE_ID_ARGS+=("$1" "$2")
      FORWARD_ARGS+=("$1" "$2")
      shift 2
      ;;
    --issue-id=*)
      ISSUE_ID_ARGS+=("$1")
      FORWARD_ARGS+=("$1")
      shift
      ;;
    *)
      FORWARD_ARGS+=("$1")
      shift
      ;;
  esac
done

CMD=(
  "${PYTHON_BIN}" "${PROJECT_ROOT}/tools/feishu_project/run_issue_data_inference.py"
  --download-root "${DOWNLOAD_ROOT}"
  --output-root "${OUTPUT_ROOT}"
  --manifest-path "${MANIFEST_PATH}"
  --python-bin "${PYTHON_BIN}"
  --inference-script "${INFERENCE_SCRIPT}"
  --video-stride "${VIDEO_STRIDE}"
)

if [[ -n "${ISSUE_JSON}" ]]; then
  CMD+=(--issue-json "${ISSUE_JSON}")
fi

if [[ "${MAX_IMAGES}" != "0" ]]; then
  CMD+=(--max-images "${MAX_IMAGES}")
fi

if [[ -n "${FRAME_INDEX_START}" ]]; then
  CMD+=(--frame-index-start "${FRAME_INDEX_START}")
fi

if [[ -n "${FRAME_INDEX_END}" ]]; then
  CMD+=(--frame-index-end "${FRAME_INDEX_END}")
fi

if [[ -n "${FRAME_ID_START}" ]]; then
  CMD+=(--frame-id-start "${FRAME_ID_START}")
fi

if [[ -n "${FRAME_ID_END}" ]]; then
  CMD+=(--frame-id-end "${FRAME_ID_END}")
fi

if [[ -n "${TARGET_FRAME_ID}" ]]; then
  CMD+=(--target-frame-id "${TARGET_FRAME_ID}")
fi

if [[ "${USE_ISSUE_FRAME_WINDOW}" != "1" && "${FRAME_BEFORE}" != "100" ]]; then
  CMD+=(--frame-before "${FRAME_BEFORE}")
fi

if [[ "${USE_ISSUE_FRAME_WINDOW}" != "1" && "${FRAME_AFTER}" != "100" ]]; then
  CMD+=(--frame-after "${FRAME_AFTER}")
fi

if [[ "${USE_ISSUE_FRAME_WINDOW}" == "1" ]]; then
  CMD+=(--use-issue-frame-window --frame-before "${FRAME_BEFORE}" --frame-after "${FRAME_AFTER}" --missing-issue-frame-policy "${MISSING_ISSUE_FRAME_POLICY}")
fi

if [[ -n "${EXPORTED_MODEL}" ]]; then
  CMD+=(--exported-model "${EXPORTED_MODEL}")
fi

if [[ -n "${DEVICE}" ]]; then
  CMD+=(--device "${DEVICE}")
fi

if [[ -n "${PROVIDERS}" ]]; then
  # shellcheck disable=SC2206
  PROVIDER_ARR=(${PROVIDERS})
  CMD+=(--providers "${PROVIDER_ARR[@]}")
fi

if [[ "${ENABLE_ATTR}" == "1" ]]; then
  CMD+=(--enable-attr)
fi

if [[ "${ENABLE_CROSS_CLASS_MERGE_PRIOR}" == "1" ]]; then
  CMD+=(--enable-cross-class-merge-prior)
fi

if [[ "${SAVE_AGGREGATE_PREDICTIONS}" == "1" ]]; then
  CMD+=(--save-aggregate-predictions)
fi

if [[ "${SKIP_EXISTING}" == "1" ]]; then
  CMD+=(--skip-existing)
fi

if [[ "${DRY_RUN}" == "1" ]]; then
  CMD+=(--dry-run)
fi

if [[ -n "${INFERENCE_EXTRA_ARGS}" ]]; then
  # shellcheck disable=SC2206
  EXTRA_ARR=(${INFERENCE_EXTRA_ARGS})
  for arg in "${EXTRA_ARR[@]}"; do
    CMD+=("--inference-arg=${arg}")
  done
fi

CMD+=("${FORWARD_ARGS[@]}")
"${CMD[@]}"

if [[ "${ENABLE_TRACKING}" != "1" ]]; then
  exit 0
fi

if [[ "${DRY_RUN}" == "1" ]]; then
  echo "Skipping tracking because DRY_RUN=1"
  exit 0
fi

TRACKING_MODEL_VERSION_RESOLVED=${TRACKING_MODEL_VERSION:-}
if [[ -z "${TRACKING_MODEL_VERSION_RESOLVED}" ]]; then
  TRACKING_MODEL_VERSION_RESOLVED=$(resolve_tracking_model_version || true)
fi

if [[ -n "${TRACKING_MODEL_VERSION_RESOLVED}" ]]; then
  echo "Running tracking with model version: ${TRACKING_MODEL_VERSION_RESOLVED}"
  TRACKING_MODEL_VERSION="${TRACKING_MODEL_VERSION_RESOLVED}" \
    PYTHON_BIN="${PYTHON_BIN}" \
    bash "${ISSUE_TRACKING_SCRIPT}" "${OUTPUT_ROOT}" "${ISSUE_ID_ARGS[@]}"
else
  echo "Running tracking without an explicit model version override"
  PYTHON_BIN="${PYTHON_BIN}" \
    bash "${ISSUE_TRACKING_SCRIPT}" "${OUTPUT_ROOT}" "${ISSUE_ID_ARGS[@]}"
fi
