cmake_minimum_required(VERSION 3.16)
project(YoloRuntime VERSION 5.0.0 LANGUAGES C CXX)
include(GNUInstallDirs)
include(CMakePackageConfigHelpers)
find_package(Python3 REQUIRED COMPONENTS Interpreter)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_C_STANDARD 11)
set(CMAKE_C_STANDARD_REQUIRED ON)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)

add_library(yolo_runtime SHARED
    src/yolo_runtime.cpp
    src/yolo_seg_geometry.cpp)
target_include_directories(yolo_runtime
    PUBLIC $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include> $<INSTALL_INTERFACE:include>
    PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
add_library(YoloRuntime::Runtime ALIAS yolo_runtime)
set_target_properties(yolo_runtime PROPERTIES EXPORT_NAME Runtime)
find_package(Threads REQUIRED)
if(WIN32)
    target_link_libraries(yolo_runtime PRIVATE Threads::Threads)
else()
    target_link_libraries(yolo_runtime PRIVATE dl Threads::Threads)
endif()
set_target_properties(yolo_runtime PROPERTIES
    VERSION 5.0.0
    SOVERSION 5
    CXX_VISIBILITY_PRESET hidden
    VISIBILITY_INLINES_HIDDEN YES)

if(WIN32)
    # Keep the public C ABI identical to Linux while exporting undecorated symbols from the DLL.
    target_sources(yolo_runtime PRIVATE yolo_runtime.def)

    set(ONNXRUNTIME_ROOT "$ENV{USERPROFILE}/.nuget/packages/microsoft.ml.onnxruntime/1.27.0"
        CACHE PATH "Root directory of Microsoft.ML.OnnxRuntime NuGet package")
    set(ONNXRUNTIME_INCLUDE_DIR "${ONNXRUNTIME_ROOT}/build/native/include")
    set(ONNXRUNTIME_LIBRARY "${ONNXRUNTIME_ROOT}/runtimes/win-x64/native/onnxruntime.lib")
    if(NOT EXISTS "${ONNXRUNTIME_INCLUDE_DIR}/onnxruntime_cxx_api.h" OR NOT EXISTS "${ONNXRUNTIME_LIBRARY}")
        message(FATAL_ERROR "ONNX Runtime 1.27.0 Windows x64 native files were not found under ${ONNXRUNTIME_ROOT}")
    endif()

    add_library(yolo_backend_onnx_cpu SHARED
        backends/onnx_cpu_backend.cpp
        yolo_backend_onnx_cpu.def)
    target_include_directories(yolo_backend_onnx_cpu PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/include
        ${CMAKE_CURRENT_SOURCE_DIR}/src
        ${ONNXRUNTIME_INCLUDE_DIR})
    target_link_libraries(yolo_backend_onnx_cpu PRIVATE ${ONNXRUNTIME_LIBRARY} Threads::Threads)
    set_property(TARGET yolo_runtime yolo_backend_onnx_cpu
        PROPERTY MSVC_RUNTIME_LIBRARY "MultiThreaded$<$<CONFIG:Debug>:Debug>")
endif()

option(YOLO_RUNTIME_BUILD_RKNN_BACKENDS "Build persistent in-process RKNN backends" ON)
option(YOLO_RUNTIME_BUILD_CONTRACT_TESTS "Compile the public C API contract test" ON)
set(RKNN_MODEL_ZOO_ROOT ""
    CACHE PATH "Path to the RKNN Model Zoo source tree")

if(YOLO_RUNTIME_BUILD_RKNN_BACKENDS AND NOT WIN32)
    if(NOT EXISTS "${RKNN_MODEL_ZOO_ROOT}/3rdparty/CMakeLists.txt")
        message(FATAL_ERROR "RKNN Model Zoo not found at ${RKNN_MODEL_ZOO_ROOT}")
    endif()

    set(TARGET_SOC rk3588 CACHE STRING "RKNN target SoC")
    add_subdirectory("${RKNN_MODEL_ZOO_ROOT}/3rdparty" model_zoo_3rdparty)
    add_subdirectory("${RKNN_MODEL_ZOO_ROOT}/utils" model_zoo_utils)

    function(add_yolo_rknn_backend target task_directory)
        add_library(${target} SHARED
            backends/rknn_backend.cpp
            src/yolo_rknn_preprocess.cpp
            ${ARGN})
        target_include_directories(${target} PRIVATE
            ${CMAKE_CURRENT_SOURCE_DIR}/include
            ${CMAKE_CURRENT_SOURCE_DIR}/src
            ${RKNN_MODEL_ZOO_ROOT}/examples/${task_directory}/cpp
            ${RKNN_MODEL_ZOO_ROOT}/utils
            ${RKNN_MODEL_ZOO_ROOT}/3rdparty/allocator/dma
            ${RKNN_MODEL_ZOO_ROOT}/3rdparty/allocator/drm
            ${LIBTIMER_INCLUDES}
            ${LIBRKNNRT_INCLUDES})
        target_link_libraries(${target} PRIVATE
            imageutils fileutils imagedrawing ${LIBRKNNRT} dl Threads::Threads)
        target_compile_definitions(${target} PRIVATE YOLO_RUNTIME_SHARED_PREPROCESS=1)
        if(target STREQUAL "yolo_backend_seg")
            set(OpenCV_DIR "${RKNN_MODEL_ZOO_ROOT}/3rdparty/opencv/opencv-linux-aarch64/share/OpenCV")
            find_package(OpenCV REQUIRED)
            target_link_libraries(${target} PRIVATE ${OpenCV_LIBS})
        endif()
        set_target_properties(${target} PROPERTIES INSTALL_RPATH "$ORIGIN")
    endfunction()

    add_yolo_rknn_backend(
        yolo_backend_detect
        yolov8
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_detect/rknpu2/yolov8.cc
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_detect/postprocess.cc)
    add_yolo_rknn_backend(
        yolo_backend_detect_yolo11
        yolov8
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_detect/rknpu2/yolov8.cc
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_detect/postprocess.cc)
    target_compile_definitions(yolo_backend_detect_yolo11 PRIVATE YOLO_MODEL_FAMILY=11)
    add_yolo_rknn_backend(
        yolo_backend_detect_yolo26
        yolov8
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_detect/rknpu2/yolov8.cc
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_detect/postprocess.cc)
    target_compile_definitions(yolo_backend_detect_yolo26 PRIVATE YOLO_MODEL_FAMILY=26)
    add_yolo_rknn_backend(
        yolo_backend_obb
        yolov8_obb
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_obb/rknpu2/yolov8-obb.cc
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_obb/postprocess.cc)
    target_compile_definitions(yolo_backend_obb PRIVATE YOLO_BACKEND_OBB=1)
    add_yolo_rknn_backend(
        yolo_backend_seg
        yolov8_seg
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_seg/rknpu2/yolov8_seg.cc
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_seg/rknpu2/postprocess.cc)
    target_compile_definitions(yolo_backend_seg PRIVATE YOLO_BACKEND_SEG=1)
    add_yolo_rknn_backend(
        yolo_backend_cls
        mobilenet
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_cls/rknpu2/mobilenet.cc)
    target_compile_definitions(yolo_backend_cls PRIVATE YOLO_BACKEND_CLS=1)
    add_yolo_rknn_backend(
        yolo_backend_pose
        yolov8_pose
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_pose/rknpu2/yolov8-pose.cc
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_pose/postprocess.cc)
    target_compile_definitions(yolo_backend_pose PRIVATE YOLO_BACKEND_POSE=1)

    if(YOLO_RUNTIME_BUILD_CONTRACT_TESTS)
        add_executable(yolo_detect_score_mapping_test
            tests/detect_score_mapping_test.cpp
            adapters/yolov8_detect/postprocess.cc)
        target_include_directories(yolo_detect_score_mapping_test PRIVATE
            ${RKNN_MODEL_ZOO_ROOT}/examples/yolov8/cpp
            ${RKNN_MODEL_ZOO_ROOT}/utils
            ${LIBRKNNRT_INCLUDES})
        add_test(NAME yolo_detect_score_mapping_test COMMAND yolo_detect_score_mapping_test)
    endif()
endif()

if(YOLO_RUNTIME_BUILD_CONTRACT_TESTS)
    enable_testing()
    add_executable(yolo_runtime_c_api_contract tests/c_api_contract.c)
    target_include_directories(yolo_runtime_c_api_contract PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)
    target_link_libraries(yolo_runtime_c_api_contract PRIVATE yolo_runtime)
    add_test(NAME yolo_runtime_c_api_contract COMMAND yolo_runtime_c_api_contract)
    add_executable(yolo_mask_bbox_test tests/mask_bbox_test.cpp)
    target_include_directories(yolo_mask_bbox_test PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/include
        ${CMAKE_CURRENT_SOURCE_DIR}/src)
    add_test(NAME yolo_mask_bbox_test COMMAND yolo_mask_bbox_test)
    add_executable(yolo_segmentation_result_filter_test tests/segmentation_result_filter_test.cpp)
    target_include_directories(yolo_segmentation_result_filter_test PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/adapters/yolov8_seg/rknpu2)
    add_test(NAME yolo_segmentation_result_filter_test COMMAND yolo_segmentation_result_filter_test)
    add_executable(yolo_windows_cpu_input_contract_test tests/windows_cpu_input_contract_test.cpp)
    target_include_directories(yolo_windows_cpu_input_contract_test PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/include
        ${CMAKE_CURRENT_SOURCE_DIR}/src)
    add_test(NAME yolo_windows_cpu_input_contract_test COMMAND yolo_windows_cpu_input_contract_test)
    add_executable(yolo_obb_nms_iou_test tests/obb_nms_iou_test.cpp)
    target_include_directories(yolo_obb_nms_iou_test PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/include
        ${CMAKE_CURRENT_SOURCE_DIR}/src)
    add_test(NAME yolo_obb_nms_iou_test COMMAND yolo_obb_nms_iou_test)
    add_executable(yolo_seg_geometry_test
        tests/seg_geometry_test.cpp
        src/yolo_seg_geometry.cpp)
    target_include_directories(yolo_seg_geometry_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
    add_test(NAME yolo_seg_geometry_test COMMAND yolo_seg_geometry_test)
    add_executable(yolo_runner_label_layout_test tests/runner_label_layout_test.cpp)
    target_include_directories(yolo_runner_label_layout_test PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/include
        ${CMAKE_CURRENT_SOURCE_DIR}/src)
    add_test(NAME yolo_runner_label_layout_test COMMAND yolo_runner_label_layout_test)
endif()

if(NOT WIN32)
    add_executable(rk3588_yolo_runner src/main.cpp)
    target_include_directories(rk3588_yolo_runner PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/include
        ${CMAKE_CURRENT_SOURCE_DIR}/src)
    target_link_libraries(rk3588_yolo_runner PRIVATE yolo_runtime)
    if(YOLO_RUNTIME_BUILD_RKNN_BACKENDS)
        target_compile_definitions(rk3588_yolo_runner PRIVATE YOLO_RUNNER_IN_PROCESS=1)
        target_include_directories(rk3588_yolo_runner PRIVATE
            ${RKNN_MODEL_ZOO_ROOT}/utils)
        target_link_libraries(rk3588_yolo_runner PRIVATE
            imageutils imagedrawing fileutils Threads::Threads)
    endif()
    set_target_properties(rk3588_yolo_runner PROPERTIES INSTALL_RPATH "$ORIGIN/lib")
endif()

if(WIN32)
    install(TARGETS yolo_runtime yolo_backend_onnx_cpu EXPORT YoloRuntimeTargets
        RUNTIME DESTINATION lib
        LIBRARY DESTINATION lib
        ARCHIVE DESTINATION lib)
    install(FILES
        "${ONNXRUNTIME_ROOT}/runtimes/win-x64/native/onnxruntime.dll"
        "${ONNXRUNTIME_ROOT}/runtimes/win-x64/native/onnxruntime_providers_shared.dll"
        DESTINATION lib)
    install(FILES "${ONNXRUNTIME_ROOT}/ThirdPartyNotices.txt" DESTINATION licenses)
else()
    install(TARGETS yolo_runtime rk3588_yolo_runner EXPORT YoloRuntimeTargets
        RUNTIME DESTINATION .
        LIBRARY DESTINATION lib)
endif()
if(YOLO_RUNTIME_BUILD_RKNN_BACKENDS AND NOT WIN32)
    install(TARGETS
        yolo_backend_detect
        yolo_backend_detect_yolo11
        yolo_backend_detect_yolo26
        yolo_backend_obb
        yolo_backend_seg
        yolo_backend_cls
        yolo_backend_pose
        LIBRARY DESTINATION lib)
endif()
install(FILES include/yolo_runtime.h include/yolo_runtime_types.h DESTINATION include)

install(EXPORT YoloRuntimeTargets NAMESPACE YoloRuntime:: DESTINATION lib/cmake/YoloRuntime)
configure_package_config_file(cmake/YoloRuntimeConfig.cmake.in
    "${CMAKE_CURRENT_BINARY_DIR}/YoloRuntimeConfig.cmake" INSTALL_DESTINATION lib/cmake/YoloRuntime)
write_basic_package_version_file("${CMAKE_CURRENT_BINARY_DIR}/YoloRuntimeConfigVersion.cmake"
    VERSION ${PROJECT_VERSION} COMPATIBILITY SameMajorVersion)
install(FILES "${CMAKE_CURRENT_BINARY_DIR}/YoloRuntimeConfig.cmake"
    "${CMAKE_CURRENT_BINARY_DIR}/YoloRuntimeConfigVersion.cmake" DESTINATION lib/cmake/YoloRuntime)
if(WIN32)
    set(CMAKE_INSTALL_SYSTEM_RUNTIME_LIBS_SKIP TRUE)
    include(InstallRequiredSystemLibraries)
    install(PROGRAMS ${CMAKE_INSTALL_SYSTEM_RUNTIME_LIBS} DESTINATION lib)
endif()
# Recompute the source identity at build time so edits never reuse stale provenance.
add_custom_target(yolo_runtime_source_identity ALL
    COMMAND "${Python3_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/scripts/artifacts.py"
        source --root "${CMAKE_CURRENT_SOURCE_DIR}" --output "${CMAKE_CURRENT_BINARY_DIR}/source-identity.json"
    VERBATIM)
add_dependencies(yolo_runtime yolo_runtime_source_identity)
install(FILES "${CMAKE_CURRENT_BINARY_DIR}/source-identity.json" DESTINATION .)
install(CODE "execute_process(COMMAND \"${Python3_EXECUTABLE}\" \"${CMAKE_CURRENT_SOURCE_DIR}/scripts/artifacts.py\" write --root \"\${CMAKE_INSTALL_PREFIX}\" RESULT_VARIABLE manifest_result)
if(NOT manifest_result EQUAL 0)
message(FATAL_ERROR \"Runtime manifest generation failed\")
endif()")

if(YOLO_RUNTIME_BUILD_CONTRACT_TESTS)
    get_property(contract_targets DIRECTORY PROPERTY BUILDSYSTEM_TARGETS)
    foreach(contract_target IN LISTS contract_targets)
        if(contract_target MATCHES "(_test|_contract)$")
            if(MSVC)
                target_compile_options(${contract_target} PRIVATE /UNDEBUG)
            else()
                target_compile_options(${contract_target} PRIVATE -UNDEBUG)
            endif()
        endif()
    endforeach()
endif()
