cmake_minimum_required(VERSION 3.18)
project(sysHAX_adapter_cpu_extension)

# 设置 C++ 标准
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

# 引入python3的cmake组件
find_package(Python3 COMPONENTS Interpreter Development.Module REQUIRED)

# 获取python拓展后缀（EXT_SUFFIX）
execute_process(
    COMMAND "${Python3_EXECUTABLE}" -c "import sysconfig; print(sysconfig.get_config_var('EXT_SUFFIX') or '')"
    OUTPUT_VARIABLE PY_EXT_SUFFIX
    OUTPUT_STRIP_TRAILING_WHITESPACE
)
if(NOT PY_EXT_SUFFIX)
    set(PY_EXT_SUFFIX ".so")
endif()

# 查找 PyTorch
find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

# 查找并连接libtorch_python.so
find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_INSTALL_PREFIX}/lib")
if(NOT TORCH_PYTHON_LIBRARY)
    message(WARNING "libtorch_python not found under ${TORCH_INSTALL_PREFIX}/lib; Python extension may fail to import.")
endif()

# 定义 define_gpu_extension_target 宏
# 这个宏用于定义 C++ 扩展目标
macro(define_gpu_extension_target EXT_NAME)
    set(options USE_SABI WITH_SOABI)
    set(oneValueArgs DESTINATION LANGUAGE)
    set(multiValueArgs SOURCES LIBRARIES COMPILE_FLAGS)
    cmake_parse_arguments(EXT "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
    
    # 创建库目标
    add_library(${EXT_NAME} SHARED ${EXT_SOURCES})

    # 确保 PYBIND11_MODULE(TORCH_EXTENSION_NAME, ...) 获取到正确的模块名称。
    target_compile_definitions(${EXT_NAME} PRIVATE TORCH_EXTENSION_NAME=${EXT_NAME})
    
    # 设置包含目录
    target_include_directories(${EXT_NAME} PRIVATE
        ${CMAKE_CURRENT_SOURCE_DIR}/csrc
        ${Python3_INCLUDE_DIRS}
        ${TORCH_INCLUDE_DIRS}
    )
    
    # 设置编译标志
    if(EXT_COMPILE_FLAGS)
        target_compile_options(${EXT_NAME} PRIVATE ${EXT_COMPILE_FLAGS})
    endif()
    
    # 链接库
    target_link_libraries(${EXT_NAME} PRIVATE ${TORCH_LIBRARIES} Python3::Module)
    if(TORCH_PYTHON_LIBRARY)
        target_link_libraries(${EXT_NAME} PRIVATE ${TORCH_PYTHON_LIBRARY})
    endif()
    if(EXT_LIBRARIES)
        target_link_libraries(${EXT_NAME} PRIVATE ${EXT_LIBRARIES})
    endif()
    
    # 设置输出目录
    if(EXT_DESTINATION)
        set_target_properties(${EXT_NAME} PROPERTIES
            LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/${EXT_DESTINATION}
        )
    endif()
    
    # 设置 Python 扩展属性
    set_target_properties(${EXT_NAME} PROPERTIES
        PREFIX ""
        SUFFIX "${PY_EXT_SUFFIX}"
    )
endmacro()

# 包含自定义的 CPU 扩展配置
include(csrc/cpu.cmake)

# ====== 支持 C++ 单元测试（gtest）======
option(BUILD_TESTS "Build C++ unit tests" OFF)
if(BUILD_TESTS)
    enable_testing()
    add_subdirectory(test/cpp)
endif()