#!/usr/bin/env bash
set -euo pipefail
shopt -s nullglob

# 用法说明：
#   适配 tools/temporal_analysis/track_objects_exported_onnx_infer_case.sh
#   的输出目录，对每个 case 下的跟踪结果文件执行协议转换。
#
#   批量模式默认会保留 case 的相对目录结构，例如：
#     输入:
#       {RESULTS_ROOT}/20251115/demo_case/merge.json
#     输出:
#       {OUTPUT_ROOT}/20251115/demo_case/ObjectPerceptionObjectList.data.json
#       {OUTPUT_ROOT}/20251115/demo_case/ObjectPerceptionObjectList.bin
#       {OUTPUT_ROOT}/20251115/demo_case/ObjectPerceptionObjectList.index.json
#   也支持直接写回每个 case 的子目录，例如：
#     输出:
#       {RESULTS_ROOT}/20251115/demo_case/{OUTPUT_DIR_NAME}/ObjectPerceptionObjectList.data.json
#       {RESULTS_ROOT}/20251115/demo_case/{OUTPUT_DIR_NAME}/ObjectPerceptionObjectList.bin
#       {RESULTS_ROOT}/20251115/demo_case/{OUTPUT_DIR_NAME}/ObjectPerceptionObjectList.index.json
#
# 支持三种模式：
#   1) 批量模式（默认）：遍历 RESULTS_ROOT 下所有 ${MERGE_JSON_NAME}
#        bash convert_merge_tracking_exported_onnx_infer_case.sh
#   2) 指定根目录模式：第一个参数传跟踪结果根目录，第二个参数可选传保存根目录
#        bash convert_merge_tracking_exported_onnx_infer_case.sh \
#             /path/to/model_20260403 \
#             /path/to/model_20260403_merge_converted
#   3) 单 case / 单文件模式：第一个参数传 case 目录、case_dir/predictions 或 merge.json
#        bash convert_merge_tracking_exported_onnx_infer_case.sh /path/to/one_case
#        bash convert_merge_tracking_exported_onnx_infer_case.sh /path/to/one_case/predictions
#        bash convert_merge_tracking_exported_onnx_infer_case.sh /path/to/one_case/merge.json
#
# 环境变量：
#   RESULTS_ROOT      默认输入根目录，和 track_objects_exported_onnx_infer_case.sh 保持一致
#   CASE_LIST_FILE    可选，批量模式下的 case 清单文件；每行一个 case 目录或 merge json 路径
#   OUTPUT_ROOT       批量模式输出根目录；OUTPUT_LAYOUT=parallel_root 时生效
#   OUTPUT_LAYOUT     批量模式输出布局：
#                       parallel_root: 输出到独立根目录并保留 case 相对路径
#                       case_subdir  : 输出到每个 case 的 {OUTPUT_DIR_NAME}/ 子目录
#   OUTPUT_DIR_NAME   单 case / case_subdir 模式默认输出目录名（默认: merge_converted）
#   MERGE_JSON_NAME   要转换的跟踪结果文件名（默认: combined_tracking.json）
#   PYTHON_BIN        Python 解释器路径
#   CAM_ID            可选，透传给 convert_merge_tracking.py 的 --cam-id

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

PYTHON_BIN="${PYTHON_BIN:-/deeplearning_team/ydong/dongying/miniconda/envs/dev/bin/python}"
RESULTS_ROOT="${RESULTS_ROOT:-/data1/dongying/Mono3d/G1Q3/model_inference/KPI/OP_KPI_TRAFFIC_SCENARIO/model_20260423-raw_no_edge}"
CASE_LIST_FILE="${CASE_LIST_FILE:-}"
OUTPUT_ROOT="${OUTPUT_ROOT:-}"
OUTPUT_LAYOUT="${OUTPUT_LAYOUT:-parallel_root}"
OUTPUT_DIR_NAME="${OUTPUT_DIR_NAME:-merge_converted}"
MERGE_JSON_NAME="${MERGE_JSON_NAME:-combined_tracking.json}"
CAM_ID="${CAM_ID:-}"
CONVERT_PARALLEL_JOBS="${CONVERT_PARALLEL_JOBS:-1}"

show_usage() {
    cat <<EOF
Usage:
  bash $(basename "${BASH_SOURCE[0]}") [TARGET_PATH] [OUTPUT_PATH]

Examples:
  bash $(basename "${BASH_SOURCE[0]}")
  bash $(basename "${BASH_SOURCE[0]}") /path/to/model_20260403
  bash $(basename "${BASH_SOURCE[0]}") /path/to/model_20260403 /path/to/model_20260403_merge_converted
  OUTPUT_ROOT=/path/to/objectlist_root \\
    bash $(basename "${BASH_SOURCE[0]}") /path/to/model_20260403
  OUTPUT_LAYOUT=case_subdir OUTPUT_DIR_NAME=objectlist \\
    bash $(basename "${BASH_SOURCE[0]}") /path/to/model_20260403
  bash $(basename "${BASH_SOURCE[0]}") /path/to/one_case
  bash $(basename "${BASH_SOURCE[0]}") /path/to/one_case /path/to/output_case_dir
  bash $(basename "${BASH_SOURCE[0]}") /path/to/one_case/combined_tracking.json /path/to/output_case_dir
EOF
}

echo ""
echo "######################################################################"
echo "# Batch convert ${MERGE_JSON_NAME} for tracked inference outputs"
echo "######################################################################"

CAM_ID_ARGS=()
if [[ -n "${CAM_ID}" ]]; then
    CAM_ID_ARGS+=(--cam-id "${CAM_ID}")
fi

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

is_predictions_dir() {
    local dir_path="$1"
    [[ -d "${dir_path}" ]] || return 1
    is_case_dir "$(dirname "${dir_path}")"
}

resolve_case_dir() {
    local target_path="$1"

    if is_case_dir "${target_path}"; then
        printf '%s\n' "${target_path}"
        return 0
    fi

    if is_predictions_dir "${target_path}"; then
        dirname "${target_path}"
        return 0
    fi

    if [[ -f "${target_path}" ]] && [[ "$(basename "${target_path}")" == "${MERGE_JSON_NAME}" ]]; then
        dirname "${target_path}"
        return 0
    fi

    return 1
}

emit_batch_merge_json_paths() {
    local batch_root="$1"

    if [[ -n "${CASE_LIST_FILE}" ]]; then
        emit_batch_merge_json_paths_from_list "${batch_root}" "${CASE_LIST_FILE}"
        return
    fi

    find "${batch_root}" -type f -name "${MERGE_JSON_NAME}" -print0 | sort -z
}

emit_batch_merge_json_paths_from_list() {
    local batch_root="${1%/}"
    local list_file="$2"
    local list_entry
    local resolved_path
    local merge_json

    if [[ ! -f "${list_file}" ]]; then
        echo "[ERROR] CASE_LIST_FILE not found: ${list_file}" >&2
        return 1
    fi

    while IFS= read -r list_entry || [[ -n "${list_entry}" ]]; do
        list_entry="${list_entry%%$'\r'}"
        [[ -n "${list_entry}" ]] || continue
        [[ "${list_entry}" == \#* ]] && continue

        if [[ "${list_entry}" == /* ]]; then
            resolved_path="${list_entry}"
        else
            resolved_path="${batch_root}/${list_entry}"
        fi

        if [[ -d "${resolved_path}" ]]; then
            merge_json="${resolved_path%/}/${MERGE_JSON_NAME}"
        else
            merge_json="${resolved_path}"
        fi

        if [[ ! -f "${merge_json}" ]]; then
            echo "[WARN] Skip missing merge json from CASE_LIST_FILE: ${resolved_path}" >&2
            continue
        fi

        printf '%s\0' "${merge_json}"
    done < "${list_file}"
}

prune_active_pids() {
    local -n pids_ref=$1
    local -a remaining=()
    local pid

    for pid in "${pids_ref[@]}"; do
        if kill -0 "${pid}" 2>/dev/null; then
            remaining+=("${pid}")
        fi
    done

    pids_ref=("${remaining[@]}")
}

default_batch_output_root() {
    local batch_root="${1%/}"
    local parent_dir
    local base_name

    parent_dir="$(dirname "${batch_root}")"
    base_name="$(basename "${batch_root}")"
    printf '%s/%s_merge_converted\n' "${parent_dir}" "${base_name}"
}

run_convert() {
    local merge_json="$1"
    local output_dir="$2"

    echo ""
    echo "=========================================="
    echo "Input  : ${merge_json}"
    echo "Output : ${output_dir}"
    echo "=========================================="

    "${PYTHON_BIN}" "${PROJECT_ROOT}/tools/convert_merge_tracking_bundle/convert_merge_tracking.py" \
        "${merge_json}" \
        -o "${output_dir}" \
        "${CAM_ID_ARGS[@]}"
}

run_single_case() {
    local case_dir="$1"
    local output_dir="$2"
    local merge_json="${case_dir}/${MERGE_JSON_NAME}"

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

    echo "Single-case mode: ${case_dir}"
    run_convert "${merge_json}" "${output_dir}"
}

run_batch_root() {
    local batch_root="${1%/}"
    local output_root="${2:-}"
    local output_layout="${3:-parallel_root}"
    local parallel_jobs="${CONVERT_PARALLEL_JOBS}"
    local total_cases=0
    local success_cases=0
    local failed_cases=0
    local status_dir
    local -a active_pids=()

    echo "Batch-root mode : ${batch_root}"
    if [[ -n "${CASE_LIST_FILE}" ]]; then
        echo "Case-list file  : ${CASE_LIST_FILE}"
    fi
    case "${output_layout}" in
        parallel_root)
            if [[ -z "${output_root}" ]]; then
                output_root="$(default_batch_output_root "${batch_root}")"
            fi
            output_root="${output_root%/}"
            echo "Output-layout   : ${output_layout}"
            echo "Output-root     : ${output_root}"
            ;;
        case_subdir)
            echo "Output-layout   : ${output_layout}"
            echo "Output-dir-name : ${OUTPUT_DIR_NAME}"
            if [[ -n "${output_root}" ]]; then
                echo "Output-root     : ${output_root} (ignored in case_subdir mode)"
            fi
            ;;
        *)
            echo "[ERROR] Unsupported OUTPUT_LAYOUT: ${output_layout}" >&2
            return 1
            ;;
    esac

    if ! [[ "${parallel_jobs}" =~ ^[0-9]+$ ]] || [[ "${parallel_jobs}" -lt 1 ]]; then
        echo "[ERROR] CONVERT_PARALLEL_JOBS must be an integer >= 1, got: ${parallel_jobs}" >&2
        return 1
    fi
    echo "Parallel jobs   : ${parallel_jobs}"

    status_dir="$(mktemp -d)"
    trap 'rm -rf "${status_dir}"' RETURN

    while IFS= read -r -d '' merge_json; do
        local case_dir
        local output_dir

        case_dir="$(dirname "${merge_json}")"
        ((total_cases+=1))

        case "${output_layout}" in
            parallel_root)
                if [[ "${case_dir}" == "${batch_root}" ]]; then
                    output_dir="${output_root}"
                else
                    output_dir="${output_root}/${case_dir#"${batch_root}/"}"
                fi
                ;;
            case_subdir)
                output_dir="${case_dir}/${OUTPUT_DIR_NAME}"
                ;;
        esac

        if [[ "${parallel_jobs}" == "1" ]]; then
            if run_convert "${merge_json}" "${output_dir}"; then
                ((success_cases+=1))
            else
                ((failed_cases+=1))
                printf '[FAIL] case=%s output=%s\n' "${case_dir}" "${output_dir}" >&2
            fi
        else
            (
                if run_convert "${merge_json}" "${output_dir}"; then
                    printf 'success\t%s\t%s\n' "${case_dir}" "${output_dir}" > "${status_dir}/$(printf '%05d' "${total_cases}").status"
                else
                    printf 'failed\t%s\t%s\n' "${case_dir}" "${output_dir}" > "${status_dir}/$(printf '%05d' "${total_cases}").status"
                    exit 1
                fi
            ) &
            active_pids+=("$!")

            while [[ "${#active_pids[@]}" -ge "${parallel_jobs}" ]]; do
                wait -n || true
                prune_active_pids active_pids
            done
        fi
    done < <(emit_batch_merge_json_paths "${batch_root}")

    if [[ "${parallel_jobs}" != "1" ]]; then
        local pid
        local status_file
        local status
        local case_dir
        local output_dir

        for pid in "${active_pids[@]}"; do
            wait "${pid}" || true
        done

        for status_file in "${status_dir}"/*.status; do
            [[ -f "${status_file}" ]] || continue
            IFS=$'\t' read -r status case_dir output_dir < "${status_file}"
            if [[ "${status}" == "success" ]]; then
                ((success_cases+=1))
            else
                ((failed_cases+=1))
                printf '[FAIL] case=%s output=%s\n' "${case_dir}" "${output_dir}" >&2
            fi
        done

        if [[ $((success_cases + failed_cases)) -ne "${total_cases}" ]]; then
            failed_cases=$((total_cases - success_cases))
        fi
    fi

    if [[ "${total_cases}" -eq 0 ]]; then
        echo "[ERROR] No ${MERGE_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 [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
    show_usage
    exit 0
fi

TARGET_PATH="${1:-${RESULTS_ROOT}}"
OUTPUT_PATH_ARG="${2:-}"

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

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