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

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/feishu_project/inference_issue_data}
DOWNLOAD_ROOT=${DOWNLOAD_ROOT:-/data1/dongying/Mono3d/G1Q3/feishu_project/downloaded_issue_data}
OUTPUT_ROOT=${OUTPUT_ROOT:-}
OUTPUT_DIR_NAME=${OUTPUT_DIR_NAME:-tracking_vis_raw}
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=""
ISSUE_IDS=()

while (($# > 0)); do
  case "$1" in
    --issue-id)
      if (($# < 2)); then
        echo "Error: --issue-id requires a value" >&2
        exit 1
      fi
      ISSUE_IDS+=("$2")
      shift 2
      ;;
    --issue-id=*)
      ISSUE_IDS+=("${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

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}")"
DOWNLOAD_ROOT_ABS="$(resolve_abs_path "${DOWNLOAD_ROOT}")"

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

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' "${DOWNLOAD_ROOT_ABS}" "${rel_case_dir}"
}

run_single_case() {
  local case_dir="$1"
  local tracking_json="${case_dir}/${TRACKING_JSON_NAME}"
  local raw_video_path
  local output_dir
  local cmd

  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}")"

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

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

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

  if [[ -n "${VIS_TRACK_IDS}" ]]; then
    # shellcheck disable=SC2206
    local track_id_arr=(${VIS_TRACK_IDS})
    cmd+=(--track-ids "${track_id_arr[@]}")
  fi

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

  echo ""
  echo "######################################################################"
  echo "# Feishu issue-data raw-image visualization"
  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 [[ "${#ISSUE_IDS[@]}" -eq 0 ]]; then
  if is_case_dir "${TARGET_PATH}"; then
    run_single_case "${TARGET_PATH}"
  elif [[ -d "${TARGET_PATH}" ]]; then
    run_batch_root "${TARGET_PATH}"
  else
    echo "Error: unsupported target path: ${TARGET_PATH}" >&2
    exit 1
  fi
  exit 0
fi

total_issues=0
success_issues=0
failed_issues=0

for issue_id in "${ISSUE_IDS[@]}"; do
  issue_root="${TARGET_PATH}/issue_${issue_id}"
  ((total_issues += 1))

  if [[ ! -d "${issue_root}" ]]; then
    echo "[FAIL] issue_${issue_id}: not found under ${TARGET_PATH}" >&2
    ((failed_issues += 1))
    continue
  fi

  if run_batch_root "${issue_root}"; then
    ((success_issues += 1))
  else
    ((failed_issues += 1))
    echo "[FAIL] issue_${issue_id}: visualization failed" >&2
  fi
done

echo ""
printf '[DONE] issues=%d success=%d failed=%d\n' \
  "${total_issues}" "${success_issues}" "${failed_issues}"

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