cmake_minimum_required(VERSION 3.12)
project(roctracer_python)

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

# 查找Python
find_package(Python3 REQUIRED COMPONENTS Interpreter Development)

# 设置路径
set(PYTHON_INCLUDE_DIR ${Python3_INCLUDE_DIRS})
set(PYBIND11_INCLUDE_DIR "/usr/local/lib/python3.10/dist-packages/torch/include")
set(HIP_INCLUDE_DIR "/opt/rocm/include")  # 根据您的HIP安装路径调整
set(ROCTRACER_INCLUDE_DIR "/opt/dtk/include")
set(ROCTRACER_LIBRARY_PATH "/opt/dtk/lib/libroctx64.so")

# 检查文件存在性
if(NOT EXISTS ${PYBIND11_INCLUDE_DIR}/pybind11/pybind11.h)
    message(FATAL_ERROR "pybind11 not found at: ${PYBIND11_INCLUDE_DIR}")
endif()

if(NOT EXISTS ${ROCTRACER_INCLUDE_DIR}/roctracer/roctx.h)
    message(FATAL_ERROR "ROCTracer headers not found at: ${ROCTRACER_INCLUDE_DIR}")
endif()

if(NOT EXISTS ${ROCTRACER_LIBRARY_PATH})
    message(FATAL_ERROR "ROCTracer library not found at: ${ROCTRACER_LIBRARY_PATH}")
endif()

# 包含目录
include_directories(
    ${CMAKE_CURRENT_SOURCE_DIR}/include
    ${PYTHON_INCLUDE_DIR}
    ${PYBIND11_INCLUDE_DIR}
    ${HIP_INCLUDE_DIR}
    ${ROCTRACER_INCLUDE_DIR}
)

# 创建C++库
add_library(roctracer_wrapper SHARED src/roctracer_wrapper.cpp)

# 添加编译定义
target_compile_definitions(roctracer_wrapper 
    PRIVATE 
        __HIP_PLATFORM_AMD__
        ROCTRACER_WRAPPER_EXPORTS
)

# 链接库
target_link_libraries(roctracer_wrapper 
    ${ROCTRACER_LIBRARY_PATH}
)

# 查找HIP库
find_library(HIP_LIBRARY hip_hcc PATHS /opt/rocm/lib)
if(HIP_LIBRARY)
    target_link_libraries(roctracer_wrapper ${HIP_LIBRARY})
endif()

# 创建Python模块
add_library(roctracer_py MODULE python/roctracer_pybind.cpp)

# 设置模块属性
set_target_properties(roctracer_py PROPERTIES
    OUTPUT_NAME "roctracer_py"
    PREFIX ""
    SUFFIX ".so"
)

# 添加编译定义
target_compile_definitions(roctracer_py 
    PRIVATE 
        __HIP_PLATFORM_AMD__
        PYBIND11_MODULE_EXPORTS
)

# 链接库
target_link_libraries(roctracer_py 
    roctracer_wrapper
    ${ROCTRACER_LIBRARY_PATH}
    ${Python3_LIBRARIES}
)

if(HIP_LIBRARY)
    target_link_libraries(roctracer_py ${HIP_LIBRARY})
endif()

# 获取Python扩展后缀
execute_process(
    COMMAND ${Python3_EXECUTABLE} -c "
import sysconfig
import sys
suffix = sysconfig.get_config_var('EXT_SUFFIX')
if suffix is None:
    suffix = sysconfig.get_config_var('SO')
sys.stdout.write(suffix if suffix else '.so')
"
    OUTPUT_VARIABLE PYTHON_MODULE_EXTENSION
    OUTPUT_STRIP_TRAILING_WHITESPACE
)

message(STATUS "Python module extension: ${PYTHON_MODULE_EXTENSION}")
set_target_properties(roctracer_py PROPERTIES SUFFIX "${PYTHON_MODULE_EXTENSION}")

# 安装
install(TARGETS roctracer_py
    LIBRARY DESTINATION lib/python${Python3_VERSION_MAJOR}.${Python3_VERSION_MINOR}/site-packages
)