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

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

CASE_LIST=${CASE_LIST:-"${PROJECT_ROOT}/tools/feishu_project/cncap_case.txt"}
SYNC_ROOT=${SYNC_ROOT:-/data1/dongying/Mono3d/CNCAP/feishu_project}
EXPORT_JSON=${EXPORT_JSON:-"${SYNC_ROOT}/exports/cncap_case_issue_list.json"}
CASE_INDEX_JSON=${CASE_INDEX_JSON:-"${SYNC_ROOT}/exports/cncap_case_issue_index.json"}
DOWNLOAD_ROOT=${DOWNLOAD_ROOT:-"${SYNC_ROOT}/downloaded_issue_data"}
DOWNLOAD_MANIFEST_PATH=${DOWNLOAD_MANIFEST_PATH:-"${DOWNLOAD_ROOT}/download_manifest.json"}
INFERENCE_ROOT=${INFERENCE_ROOT:-"${SYNC_ROOT}/inference_issue_data"}
INFERENCE_MANIFEST_PATH=${INFERENCE_MANIFEST_PATH:-"${INFERENCE_ROOT}/inference_manifest.json"}

PYTHON_BIN=${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}
RUN_BUILD=${RUN_BUILD:-1}
RUN_DOWNLOAD=${RUN_DOWNLOAD:-1}
RUN_INFERENCE=${RUN_INFERENCE:-1}
ENABLE_TRACKING=${ENABLE_TRACKING:-1}
ENABLE_CONVERT=${ENABLE_CONVERT:-1}
ENABLE_VISUALIZE=${ENABLE_VISUALIZE:-0}
ENABLE_TEMPORAL_ANALYSIS=${ENABLE_TEMPORAL_ANALYSIS:-0}
SKIP_EXISTING_INFERENCE=${SKIP_EXISTING_INFERENCE:-1}
DRY_RUN=${DRY_RUN:-0}

USE_ISSUE_FRAME_WINDOW=${USE_ISSUE_FRAME_WINDOW:-0}
FRAME_BEFORE=${FRAME_BEFORE:-200}
FRAME_AFTER=${FRAME_AFTER:-200}
MISSING_ISSUE_FRAME_POLICY=${MISSING_ISSUE_FRAME_POLICY:-full}
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:-}

VIDEO_STRIDE=${VIDEO_STRIDE:-1}
MAX_IMAGES=${MAX_IMAGES:-0}
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}
INFERENCE_EXTRA_ARGS=${INFERENCE_EXTRA_ARGS:-}

TRACK_CLASSES=${TRACK_CLASSES:-0 1 2 3 4 5 6 7 8 9 10 11 12}
IOU_THRESH=${IOU_THRESH:-0.3}
MAX_AGE=${MAX_AGE:-5}
MIN_HITS=${MIN_HITS:-1}
DIST_THRESH=${DIST_THRESH:-100}
MAX_3D_DISTANCE=${MAX_3D_DISTANCE:-10.0}
MAX_FRAMES=${MAX_FRAMES:-}
MERGE_OUTPUT_NAME=${MERGE_OUTPUT_NAME:-combined_tracking.json}
FILE_PATTERN=${FILE_PATTERN:-*.json}
ENABLE_USE_3D=${ENABLE_USE_3D:-0}
CONVERT_OUTPUT_DIR_NAME=${CONVERT_OUTPUT_DIR_NAME:-objectlist}
CONVERT_TRACKING_JSON_NAME=${CONVERT_TRACKING_JSON_NAME:-merge.json}
CONVERT_CAM_ID=${CONVERT_CAM_ID:-}
VIS_DOWNLOAD_ROOT=${VIS_DOWNLOAD_ROOT:-"${DOWNLOAD_ROOT}"}
VIS_OUTPUT_ROOT=${VIS_OUTPUT_ROOT:-}
VIS_OUTPUT_DIR_NAME=${VIS_OUTPUT_DIR_NAME:-tracking_vis_raw}
VIS_TRACKING_JSON_NAME=${VIS_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}
TEMPORAL_OUTPUT_ROOT=${TEMPORAL_OUTPUT_ROOT:-}
TEMPORAL_OUTPUT_DIR_NAME=${TEMPORAL_OUTPUT_DIR_NAME:-temporal_observation}
TEMPORAL_JSON_NAME=${TEMPORAL_JSON_NAME:-merge.json}
TEMPORAL_MIN_LENGTH=${TEMPORAL_MIN_LENGTH:-3}
TEMPORAL_TRACK_IDS=${TEMPORAL_TRACK_IDS:-}
TEMPORAL_CLASS_ID=${TEMPORAL_CLASS_ID:-}
TEMPORAL_FRAME_ID_START=${TEMPORAL_FRAME_ID_START:-}
TEMPORAL_FRAME_ID_END=${TEMPORAL_FRAME_ID_END:-}
TEMPORAL_PLOTS=${TEMPORAL_PLOTS:-0}
TEMPORAL_EXPORT_SERIES=${TEMPORAL_EXPORT_SERIES:-1}
TEMPORAL_FOCUS_TRACK_PLOTS=${TEMPORAL_FOCUS_TRACK_PLOTS:-1}
TEMPORAL_PREFER_EGO=${TEMPORAL_PREFER_EGO:-1}
TEMPORAL_X_AXIS=${TEMPORAL_X_AXIS:-frame_id}
TEMPORAL_HEADING_SOURCE=${TEMPORAL_HEADING_SOURCE:-camera_reg}
TRACKING_MODEL_VERSION=${TRACKING_MODEL_VERSION:-}

SKIP_MDI=${SKIP_MDI:-0}
SKIP_COPY=${SKIP_COPY:-0}
ONLY_REDOWNLOAD_AFFECTED_CASES=${ONLY_REDOWNLOAD_AFFECTED_CASES:-0}
ID_BASE=${ID_BASE:-9000000000}

CASE_JSON_BUILDER=${CASE_JSON_BUILDER:-"${PROJECT_ROOT}/tools/feishu_project/build_cncap_case_issue_json.py"}
DOWNLOAD_WRAPPER=${DOWNLOAD_WRAPPER:-"${PROJECT_ROOT}/tools/feishu_project/download_issue_data.sh"}
INFERENCE_WRAPPER=${INFERENCE_WRAPPER:-"${PROJECT_ROOT}/tools/feishu_project/run_issue_data_inference.sh"}
TRACKING_WRAPPER=${TRACKING_WRAPPER:-"${PROJECT_ROOT}/tools/feishu_project/run_issue_data_tracking.sh"}

ISSUE_ID_ARGS=()
ISSUE_ID_VALUES=()

show_usage() {
  cat <<EOF
Usage:
  bash $(basename "${BASH_SOURCE[0]}") [--issue-id <id>]...

Environment variables:
  CASE_LIST, SYNC_ROOT, EXPORT_JSON, CASE_INDEX_JSON
  DOWNLOAD_ROOT, INFERENCE_ROOT, EXPORTED_MODEL
  RUN_BUILD, RUN_DOWNLOAD, RUN_INFERENCE, ENABLE_TRACKING
  ENABLE_CONVERT, ENABLE_VISUALIZE, ENABLE_TEMPORAL_ANALYSIS
  DRY_RUN, ID_BASE
EOF
}

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")
      ISSUE_ID_VALUES+=("$2")
      shift 2
      ;;
    --issue-id=*)
      issue_id="${1#*=}"
      ISSUE_ID_ARGS+=("$1")
      ISSUE_ID_VALUES+=("${issue_id}")
      shift
      ;;
    -h|--help)
      show_usage
      exit 0
      ;;
    *)
      echo "Error: unsupported option: $1" >&2
      show_usage >&2
      exit 1
      ;;
  esac
done

print_stage() {
  local stage_name="$1"
  echo ""
  echo "######################################################################"
  echo "# ${stage_name}"
  echo "######################################################################"
}

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
}

has_local_video_cases() {
  local search_root="$1"
  shift || true

  if [[ ! -d "${search_root}" ]]; then
    return 1
  fi

  if [[ "$#" -eq 0 ]]; then
    [[ -n "$(find "${search_root}" -type f -path '*/sigmastar.1/camera4.bin' -print -quit)" ]]
    return
  fi

  local issue_id=""
  for issue_id in "$@"; do
    if [[ -d "${search_root}/issue_${issue_id}" ]] && [[ -n "$(find "${search_root}/issue_${issue_id}" -type f -path '*/sigmastar.1/camera4.bin' -print -quit)" ]]; then
      return 0
    fi
  done
  return 1
}

TMP_ROOT=""
EFFECTIVE_EXPORT_JSON="${EXPORT_JSON}"
EFFECTIVE_CASE_INDEX_JSON="${CASE_INDEX_JSON}"

cleanup() {
  if [[ -n "${TMP_ROOT}" && -d "${TMP_ROOT}" ]]; then
    rm -rf "${TMP_ROOT}"
  fi
}
trap cleanup EXIT

if [[ "${DRY_RUN}" == "1" && "${RUN_BUILD}" == "1" ]]; then
  TMP_ROOT=$(mktemp -d /tmp/cncap_case_pipeline.XXXXXX)
  EFFECTIVE_EXPORT_JSON="${TMP_ROOT}/cncap_case_issue_list.json"
  EFFECTIVE_CASE_INDEX_JSON="${TMP_ROOT}/cncap_case_issue_index.json"
fi

export PYTHON_BIN
export TRACK_CLASSES
export IOU_THRESH
export MAX_AGE
export MIN_HITS
export DIST_THRESH
export MAX_3D_DISTANCE
export MAX_FRAMES
export MERGE_OUTPUT_NAME
export FILE_PATTERN
export ENABLE_USE_3D
export ENABLE_CONVERT
export ENABLE_VISUALIZE
export ENABLE_TEMPORAL_ANALYSIS
export CONVERT_OUTPUT_DIR_NAME
export CONVERT_TRACKING_JSON_NAME
export CONVERT_CAM_ID
export VIS_DOWNLOAD_ROOT
export VIS_OUTPUT_ROOT
export VIS_OUTPUT_DIR_NAME
export VIS_TRACKING_JSON_NAME
export VIS_MAX_FRAMES
export VIS_CLASS_ID
export VIS_TRACK_IDS
export VIS_SHOW_TRAJECTORY
export TEMPORAL_OUTPUT_ROOT
export TEMPORAL_OUTPUT_DIR_NAME
export TEMPORAL_JSON_NAME
export TEMPORAL_MIN_LENGTH
export TEMPORAL_TRACK_IDS
export TEMPORAL_CLASS_ID
export TEMPORAL_FRAME_ID_START
export TEMPORAL_FRAME_ID_END
export TEMPORAL_PLOTS
export TEMPORAL_EXPORT_SERIES
export TEMPORAL_FOCUS_TRACK_PLOTS
export TEMPORAL_PREFER_EGO
export TEMPORAL_X_AXIS
export TEMPORAL_HEADING_SOURCE

if [[ "${RUN_BUILD}" == "1" ]]; then
  print_stage "Build CNCAP Synthetic Issue JSON"
  echo "CASE_LIST         : ${CASE_LIST}"
  echo "EXPORT_JSON       : ${EFFECTIVE_EXPORT_JSON}"
  echo "CASE_INDEX_JSON   : ${EFFECTIVE_CASE_INDEX_JSON}"
  echo "ID_BASE           : ${ID_BASE}"

  "${PYTHON_BIN}" "${CASE_JSON_BUILDER}" \
    --input "${CASE_LIST}" \
    --output "${EFFECTIVE_EXPORT_JSON}" \
    --case-index-output "${EFFECTIVE_CASE_INDEX_JSON}" \
    --id-base "${ID_BASE}"
elif [[ ! -f "${EFFECTIVE_EXPORT_JSON}" ]]; then
  echo "Error: export json does not exist and RUN_BUILD=0: ${EFFECTIVE_EXPORT_JSON}" >&2
  exit 1
fi

if [[ "${RUN_DOWNLOAD}" == "1" ]]; then
  print_stage "Download CNCAP Case Data"
  INPUT_JSON="${EFFECTIVE_EXPORT_JSON}" \
    OUTPUT_ROOT="${DOWNLOAD_ROOT}" \
    MANIFEST_PATH="${DOWNLOAD_MANIFEST_PATH}" \
    PYTHON_BIN="${PYTHON_BIN}" \
    DRY_RUN="${DRY_RUN}" \
    SKIP_MDI="${SKIP_MDI}" \
    SKIP_COPY="${SKIP_COPY}" \
    ONLY_REDOWNLOAD_AFFECTED_CASES="${ONLY_REDOWNLOAD_AFFECTED_CASES}" \
    bash "${DOWNLOAD_WRAPPER}" "${ISSUE_ID_ARGS[@]}"
fi

if [[ "${RUN_INFERENCE}" == "1" ]]; then
  if [[ "${DRY_RUN}" == "1" ]] && ! has_local_video_cases "${DOWNLOAD_ROOT}" "${ISSUE_ID_VALUES[@]}"; then
    print_stage "Skip CNCAP Inference"
    echo "Skipping inference in dry-run because no local downloaded camera4.bin cases were found."
    echo "Run download without DRY_RUN first, or reuse an existing DOWNLOAD_ROOT to plan inference."
  else
    print_stage "Run CNCAP Inference"
    DOWNLOAD_ROOT="${DOWNLOAD_ROOT}" \
      OUTPUT_ROOT="${INFERENCE_ROOT}" \
      MANIFEST_PATH="${INFERENCE_MANIFEST_PATH}" \
      PYTHON_BIN="${PYTHON_BIN}" \
      ISSUE_JSON="${EFFECTIVE_EXPORT_JSON}" \
      VIDEO_STRIDE="${VIDEO_STRIDE}" \
      MAX_IMAGES="${MAX_IMAGES}" \
      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}" \
      FRAME_AFTER="${FRAME_AFTER}" \
      USE_ISSUE_FRAME_WINDOW="${USE_ISSUE_FRAME_WINDOW}" \
      MISSING_ISSUE_FRAME_POLICY="${MISSING_ISSUE_FRAME_POLICY}" \
      EXPORTED_MODEL="${EXPORTED_MODEL}" \
      DEVICE="${DEVICE}" \
      PROVIDERS="${PROVIDERS}" \
      ENABLE_ATTR="${ENABLE_ATTR}" \
      ENABLE_CROSS_CLASS_MERGE_PRIOR="${ENABLE_CROSS_CLASS_MERGE_PRIOR}" \
      SAVE_AGGREGATE_PREDICTIONS="${SAVE_AGGREGATE_PREDICTIONS}" \
      SKIP_EXISTING="${SKIP_EXISTING_INFERENCE}" \
      DRY_RUN="${DRY_RUN}" \
      ENABLE_TRACKING=0 \
      INFERENCE_EXTRA_ARGS="${INFERENCE_EXTRA_ARGS}" \
      bash "${INFERENCE_WRAPPER}" "${ISSUE_ID_ARGS[@]}"
  fi
fi

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

if [[ "${DRY_RUN}" == "1" ]]; then
  print_stage "Skip CNCAP Tracking"
  echo "Skipping tracking because DRY_RUN=1."
  exit 0
fi

TRACKING_MODEL_VERSION_RESOLVED=$(resolve_tracking_model_version || true)

print_stage "Run CNCAP Tracking/Convert Workflow"
echo "INFERENCE_ROOT          : ${INFERENCE_ROOT}"
echo "VIS_DOWNLOAD_ROOT       : ${VIS_DOWNLOAD_ROOT}"
if [[ -n "${TRACKING_MODEL_VERSION_RESOLVED}" ]]; then
  echo "TRACKING_MODEL_VERSION  : ${TRACKING_MODEL_VERSION_RESOLVED}"
fi

RESULTS_ROOT="${INFERENCE_ROOT}" \
  PYTHON_BIN="${PYTHON_BIN}" \
  TRACKING_MODEL_VERSION="${TRACKING_MODEL_VERSION_RESOLVED}" \
  bash "${TRACKING_WRAPPER}" "${INFERENCE_ROOT}" "${ISSUE_ID_ARGS[@]}"