include(${CMAKE_SOURCE_DIR}/cmake/cxx_standard.cmake)

# Add an option to enable/disable CUDA kernel tests. This option is ON by default
# if building on non-Windows platform with CUDA available, and OFF otherwise.
cmake_dependent_option(ENABLE_CUDA_KERNEL_TESTS "Build cuda kernel tests" ON "USE_CUDA;CMAKE_CUDA_COMPILER" OFF)

# unit tests program
add_executable(unit_tests)

file(GLOB test_srcs CONFIGURE_DEPENDS
  "${CMAKE_CURRENT_SOURCE_DIR}/*.h"
  "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp"
)

# Keep tests with dedicated white-box executables out of the public-API unit_tests target.
list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/reinit_tests.cpp")
list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/search_checkpoint_tests.cpp")

if(USE_CUDA AND CMAKE_CUDA_COMPILER AND ENABLE_CUDA_KERNEL_TESTS)
  message(STATUS "Including CUDA kernel tests in the build.")
  file(GLOB cuda_kernel_test_srcs CONFIGURE_DEPENDS
    "${CMAKE_CURRENT_SOURCE_DIR}/cuda_kernel/*.h"
    "${CMAKE_CURRENT_SOURCE_DIR}/cuda_kernel/*.cpp"
  )
  target_sources(unit_tests PRIVATE ${test_srcs} ${generator_cudalib_srcs} ${cuda_kernel_test_srcs})
  # Enable STABLE_TOPK in Windows so that we have test coverage for stable sort.
  if (CMAKE_SYSTEM_NAME STREQUAL "Windows")
    message(STATUS "Enable STABLE_TOPK in CUDA kernel tests.")
    target_compile_definitions(unit_tests PRIVATE STABLE_TOPK)
  endif()
else()
  target_sources(unit_tests PRIVATE ${test_srcs})
endif()

target_include_directories(unit_tests PRIVATE
  ${ORT_HEADER_DIR}
  ${onnxruntime_extensions_SOURCE_DIR}/shared/api
  ${CMAKE_SOURCE_DIR}/src
)

target_link_directories(unit_tests PRIVATE ${ORT_LIB_DIR})
target_link_libraries(unit_tests PRIVATE
  onnxruntime-genai
  onnxruntime_extensions
  GTest::gtest
)
set_target_properties(unit_tests PROPERTIES
    RUNTIME_OUTPUT_DIRECTORY "$<TARGET_FILE_DIR:onnxruntime-genai>"
)

if(NOT (CMAKE_SYSTEM_NAME STREQUAL "Android" OR CMAKE_SYSTEM_NAME STREQUAL "Linux" OR CMAKE_SYSTEM_NAME STREQUAL "Darwin"))
  target_link_libraries(unit_tests PRIVATE ${ONNXRUNTIME_LIB})
endif()

if(USE_CUDA AND CMAKE_CUDA_COMPILER AND ENABLE_CUDA_KERNEL_TESTS)
  target_link_libraries(unit_tests PRIVATE cublasLt cublas curand cufft cudart)
  set_target_properties(unit_tests PROPERTIES LINKER_LANGUAGE CUDA)
  add_dependencies(unit_tests onnxruntime-genai-cuda)

  # These tests exercise the internal BatchedSampler interface, so keep them out of unit_tests,
  # which deliberately links only the public shared-library API.
  add_executable(cuda_batched_sampler_tests
    "${CMAKE_CURRENT_SOURCE_DIR}/cuda_batched_sampler_tests/cuda_batched_sampler_tests.cpp")
  target_include_directories(cuda_batched_sampler_tests PRIVATE
    ${ORT_HEADER_DIR}
    ${onnxruntime_extensions_SOURCE_DIR}/shared/api
    ${CMAKE_SOURCE_DIR}/src
    ${CMAKE_CURRENT_SOURCE_DIR}
  )
  target_link_libraries(cuda_batched_sampler_tests PRIVATE
    onnxruntime-genai-obj
    GTest::gtest
  )
  add_dependencies(cuda_batched_sampler_tests onnxruntime-genai)
  set_target_properties(cuda_batched_sampler_tests PROPERTIES
    RUNTIME_OUTPUT_DIRECTORY "$<TARGET_FILE_DIR:onnxruntime-genai>"
    FOLDER "Tests"
  )
  if(NOT MSVC)
    target_compile_options(cuda_batched_sampler_tests PRIVATE "-fvisibility=hidden")
  endif()
  add_test(NAME CudaBatchedSamplerTests COMMAND cuda_batched_sampler_tests)
endif()

set(TEST_MODEL_SRC_DIR "${CMAKE_CURRENT_SOURCE_DIR}/models/")

add_compile_definitions(MODEL_PATH="${TEST_MODEL_SRC_DIR}")
set_target_properties(unit_tests PROPERTIES FOLDER "Tests")
get_target_property(all_test_srcs unit_tests SOURCES)
source_group(TREE ${PROJECT_SOURCE_DIR} FILES ${all_test_srcs})
set_property(DIRECTORY ${CMAKE_SOURCE_DIR} PROPERTY VS_STARTUP_PROJECT unit_tests)

# Hide symbols by default, so that shared libraries don't link to our redirected symbols (leads to infinite loops)
if (NOT MSVC)
  target_compile_options(unit_tests PRIVATE "-fvisibility=hidden")
endif()

add_test(NAME UnitTests COMMAND unit_tests)

# Standalone shutdown / re-initialization tests. Isolated in their own executable because they call
# OgaShutdown(), which resets process-global GenAI state (env, device interfaces, add-on libraries,
# registered plugin EPs). Keeping them out of unit_tests avoids any coupling to test order. This
# binary links the genai OBJECT library (not the shared DLL) so it can call internal genai symbols
# (e.g. GetDeviceInterface) directly for white-box coverage; unit_tests deliberately stays on the
# public API via the shared library.
add_executable(reinit_tests "${CMAKE_CURRENT_SOURCE_DIR}/reinit_tests.cpp")
target_include_directories(reinit_tests PRIVATE
  ${ORT_HEADER_DIR}
  ${onnxruntime_extensions_SOURCE_DIR}/shared/api
  ${CMAKE_SOURCE_DIR}/src
)
# onnxruntime-genai-obj provides genai's objects and its PUBLIC usage requirements (ORT + extensions
# include dirs; extensions / ORT / Threads link libs), so those are inherited here.
target_link_libraries(reinit_tests PRIVATE
  onnxruntime-genai-obj
  GTest::gtest
)
# The shared library stages the ORT runtime (onnxruntime.dll) and the CUDA add-on next to itself;
# depend on it so those are present in reinit_tests' output directory at run time.
add_dependencies(reinit_tests onnxruntime-genai)
set_target_properties(reinit_tests PROPERTIES
    RUNTIME_OUTPUT_DIRECTORY "$<TARGET_FILE_DIR:onnxruntime-genai>"
    FOLDER "Tests"
)
if (NOT MSVC)
  target_compile_options(reinit_tests PRIVATE "-fvisibility=hidden")
endif()

add_test(NAME ReInitTests COMMAND reinit_tests)

if (ENABLE_TELEMETRY AND NOT WIN32 AND NOT ANDROID AND NOT CMAKE_SYSTEM_NAME STREQUAL "iOS")
  add_executable(telemetry_device_info_tests
    "${CMAKE_CURRENT_SOURCE_DIR}/telemetry/device_info_tests.cpp")
  target_include_directories(telemetry_device_info_tests PRIVATE
    ${ORT_HEADER_DIR}
    ${onnxruntime_extensions_SOURCE_DIR}/shared/api
    ${CMAKE_SOURCE_DIR}/src
    ${CMAKE_CURRENT_SOURCE_DIR}
  )
  target_link_libraries(telemetry_device_info_tests PRIVATE
    onnxruntime-genai-obj
    GTest::gtest
  )
  add_dependencies(telemetry_device_info_tests onnxruntime-genai)
  set_target_properties(telemetry_device_info_tests PROPERTIES
    RUNTIME_OUTPUT_DIRECTORY "$<TARGET_FILE_DIR:onnxruntime-genai>"
    FOLDER "Tests"
  )
  if(NOT MSVC)
    target_compile_options(telemetry_device_info_tests PRIVATE "-fvisibility=hidden")
  endif()
  add_test(NAME TelemetryDeviceInfoTests COMMAND telemetry_device_info_tests)
endif()

# Dedicated white-box unit tests for the continuous-batching Engine. Like reinit_tests, this binary
# links the genai OBJECT library (not the public shared DLL) so it can exercise internal Engine
# types (Block/BlockPool, the scheduler, cache manager, model executor, and Request invariants)
# directly. It is kept separate from unit_tests so that internal test-only fakes/helpers do not leak
# into the public-API suite and Engine failures are attributable to Engine correctness.
# It hosts pure unit tests (no model, no ONNX session, no GPU) alongside component
# tests that mint requests from a tiny CPU fixture model; the fixture model only loads and never
# runs the graph, so the suite still needs no GPU and no real inference.
file(GLOB engine_unit_test_srcs CONFIGURE_DEPENDS
  "${CMAKE_CURRENT_SOURCE_DIR}/engine/*.h"
  "${CMAKE_CURRENT_SOURCE_DIR}/engine/*.cpp"
)
list(APPEND engine_unit_test_srcs
  "${CMAKE_CURRENT_SOURCE_DIR}/search_checkpoint_tests.cpp"
)
add_executable(engine_unit_tests ${engine_unit_test_srcs})
target_include_directories(engine_unit_tests PRIVATE
  ${ORT_HEADER_DIR}
  ${onnxruntime_extensions_SOURCE_DIR}/shared/api
  ${CMAKE_SOURCE_DIR}/src
  ${CMAKE_CURRENT_SOURCE_DIR}
)
# onnxruntime-genai-obj propagates its PUBLIC usage requirements (ORT + extensions include dirs and
# link libs), so those are inherited here.
target_link_libraries(engine_unit_tests PRIVATE
  onnxruntime-genai-obj
  GTest::gtest
)
# The shared library stages the ORT runtime (and CUDA add-on) next to itself; depend on it so those
# are present in engine_unit_tests' output directory at run time.
add_dependencies(engine_unit_tests onnxruntime-genai)
set_target_properties(engine_unit_tests PROPERTIES
    RUNTIME_OUTPUT_DIRECTORY "$<TARGET_FILE_DIR:onnxruntime-genai>"
    FOLDER "Tests"
)
if (NOT MSVC)
  target_compile_options(engine_unit_tests PRIVATE "-fvisibility=hidden")
endif()

add_test(NAME EngineUnitTests COMMAND engine_unit_tests)
set_tests_properties(EngineUnitTests PROPERTIES LABELS "engine;unit")
