cmake_minimum_required(VERSION 3.15...3.27)
project(qkmeans VERSION 0.1.0 LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_POSITION_INDEPENDENT_CODE ON)

# Options - defaults optimized for portable pre-built wheels
option(QKMEANS_USE_OPENMP "Enable OpenMP parallelization" OFF)
option(QKMEANS_USE_AVX2 "Enable AVX2 SIMD optimizations" OFF)
option(QKMEANS_USE_FAISS "Enable FAISS backend (if available)" OFF)
option(QKMEANS_USE_HNSWLIB "Enable hnswlib backend (if available)" OFF)

# Find dependencies
find_package(pybind11 CONFIG REQUIRED)

# OpenMP (tricky on macOS, disabled by default for portability)
if(QKMEANS_USE_OPENMP)
    if(APPLE)
        # macOS: Try Homebrew's libomp
        execute_process(
            COMMAND brew --prefix libomp
            OUTPUT_VARIABLE HOMEBREW_LIBOMP_PREFIX
            OUTPUT_STRIP_TRAILING_WHITESPACE
            ERROR_QUIET
        )
        if(HOMEBREW_LIBOMP_PREFIX)
            set(OpenMP_CXX_FLAGS "-Xpreprocessor -fopenmp -I${HOMEBREW_LIBOMP_PREFIX}/include")
            set(OpenMP_CXX_LIB_NAMES "omp")
            set(OpenMP_omp_LIBRARY "${HOMEBREW_LIBOMP_PREFIX}/lib/libomp.dylib")
        endif()
    endif()
    find_package(OpenMP)
    if(OpenMP_CXX_FOUND)
        message(STATUS "OpenMP found: ${OpenMP_CXX_VERSION}")
    else()
        message(WARNING "OpenMP not found, disabling parallel execution")
        set(QKMEANS_USE_OPENMP OFF)
    endif()
endif()

# SIMD flags (disabled by default for portable wheels)
set(QKMEANS_SIMD_FLAGS "")
if(QKMEANS_USE_AVX2)
    include(CheckCXXCompilerFlag)
    if(MSVC)
        check_cxx_compiler_flag("/arch:AVX2" COMPILER_SUPPORTS_AVX2)
        if(COMPILER_SUPPORTS_AVX2)
            set(QKMEANS_SIMD_FLAGS "/arch:AVX2")
        endif()
    else()
        check_cxx_compiler_flag("-mavx2" COMPILER_SUPPORTS_AVX2)
        if(COMPILER_SUPPORTS_AVX2)
            set(QKMEANS_SIMD_FLAGS "-mavx2 -mfma")
        endif()
    endif()
    if(COMPILER_SUPPORTS_AVX2)
        add_compile_definitions(QKMEANS_HAS_AVX2)
        message(STATUS "AVX2 support enabled")
    else()
        message(WARNING "AVX2 not supported by compiler")
    endif()
endif()

# Optional: FAISS
if(QKMEANS_USE_FAISS)
    find_package(faiss QUIET)
    if(faiss_FOUND)
        add_compile_definitions(QKMEANS_HAS_FAISS)
        message(STATUS "FAISS found")
    else()
        message(WARNING "FAISS not found, disabling FAISS backend")
        set(QKMEANS_USE_FAISS OFF)
    endif()
endif()

# Optional: hnswlib (header-only)
if(QKMEANS_USE_HNSWLIB)
    if(NOT HNSWLIB_INCLUDE_DIR)
        find_path(HNSWLIB_INCLUDE_DIR hnswlib/hnswlib.h)
    endif()
    if(HNSWLIB_INCLUDE_DIR)
        add_compile_definitions(QKMEANS_HAS_HNSWLIB)
        message(STATUS "hnswlib found at ${HNSWLIB_INCLUDE_DIR}")
    else()
        message(WARNING "hnswlib not found, disabling HNSW backend")
        set(QKMEANS_USE_HNSWLIB OFF)
    endif()
endif()

# Source files
set(QKMEANS_SOURCES
    src/distance.cpp
    src/sampling.cpp
    src/anns/lsh.cpp
    src/anns/brute_force.cpp
    src/core.cpp
    src/bindings.cpp
)

if(QKMEANS_USE_FAISS)
    list(APPEND QKMEANS_SOURCES src/anns/faiss_backend.cpp)
endif()

if(QKMEANS_USE_HNSWLIB)
    list(APPEND QKMEANS_SOURCES src/anns/hnsw_backend.cpp)
endif()

# Create the Python module
pybind11_add_module(_qkmeans_cpp ${QKMEANS_SOURCES})

target_include_directories(_qkmeans_cpp PRIVATE
    ${CMAKE_CURRENT_SOURCE_DIR}/include
)

# Apply SIMD flags
if(QKMEANS_SIMD_FLAGS)
    target_compile_options(_qkmeans_cpp PRIVATE ${QKMEANS_SIMD_FLAGS})
endif()

# Link OpenMP
if(QKMEANS_USE_OPENMP AND OpenMP_CXX_FOUND)
    target_link_libraries(_qkmeans_cpp PRIVATE OpenMP::OpenMP_CXX)
endif()

# Link FAISS
if(QKMEANS_USE_FAISS AND faiss_FOUND)
    target_link_libraries(_qkmeans_cpp PRIVATE faiss)
endif()

# Include hnswlib
if(QKMEANS_USE_HNSWLIB AND HNSWLIB_INCLUDE_DIR)
    target_include_directories(_qkmeans_cpp PRIVATE ${HNSWLIB_INCLUDE_DIR})
endif()

# Install the module
install(TARGETS _qkmeans_cpp LIBRARY DESTINATION qkmeans)
