cmake_minimum_required(VERSION 3.18)
project(mlipcpp VERSION 0.1.2 LANGUAGES CXX C)

set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)

set(CMAKE_C_STANDARD 11)
set(CMAKE_C_STANDARD_REQUIRED ON)

# Set output directories (can be overridden by user)
if(NOT CMAKE_RUNTIME_OUTPUT_DIRECTORY)
    set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
endif()
if(NOT CMAKE_LIBRARY_OUTPUT_DIRECTORY)
    set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
endif()
if(NOT CMAKE_ARCHIVE_OUTPUT_DIRECTORY)
    set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/lib)
endif()

# =============================================================================
# Options
# =============================================================================
option(MLIPCPP_BUILD_EXAMPLES "Build examples" ON)
option(MLIPCPP_BUILD_TESTS "Build tests" OFF)
option(MLIPCPP_BUILD_FORTRAN "Build Fortran bindings" OFF)
option(MLIPCPP_BUILD_PYTHON "Build Python bindings" OFF)
option(MLIPCPP_INSTALL "Generate install target" OFF)
option(MLIPCPP_USE_CUDA "Enable CUDA backend via ggml" OFF)
option(MLIPCPP_USE_HIP "Enable HIP/ROCm backend via ggml (AMD)" OFF)
option(MLIPCPP_USE_METAL "Enable Metal backend via ggml" OFF)
option(MLIPCPP_USE_WEBGPU "Enable WebGPU backend via ggml (requires Dawn)" OFF)
option(MLIPCPP_USE_VULKAN "Enable Vulkan backend via ggml" OFF)
option(MLIPCPP_USE_SYCL "Enable SYCL backend via ggml (Intel)" OFF)
option(MLIPCPP_USE_CANN "Enable CANN backend via ggml (Huawei Ascend)" OFF)
option(MLIPCPP_USE_BLAS "Enable BLAS acceleration via ggml" OFF)
option(MLIPCPP_USE_SYSTEM_FMT "Use system-installed fmtlib instead of bundled" OFF)
option(MLIPCPP_WASM_ASYNCIFY "Use ASYNCIFY instead of JSPI for async WebGPU calls in WASM builds" OFF)

# WASM-specific settings
if(EMSCRIPTEN)
    set(MLIPCPP_BUILD_WASM ON)
    set(MLIPCPP_BUILD_EXAMPLES OFF)
    set(MLIPCPP_BUILD_TESTS OFF)
    set(MLIPCPP_BUILD_PYTHON OFF)
    set(MLIPCPP_BUILD_FORTRAN OFF)
    # Disable native GPU backends for WASM (WebGPU is the only option)
    set(MLIPCPP_USE_CUDA OFF)
    set(MLIPCPP_USE_HIP OFF)
    set(MLIPCPP_USE_METAL OFF)
    set(MLIPCPP_USE_VULKAN OFF)
    set(MLIPCPP_USE_SYCL OFF)
    set(MLIPCPP_USE_CANN OFF)
    set(MLIPCPP_USE_BLAS OFF)
    # MLIPCPP_USE_WEBGPU stays as the user set it (default OFF).
    # Tell GGML we're building for WASM so it uses SIMD-optimized kernels
    # (Emscripten defaults CMAKE_SYSTEM_PROCESSOR to x86, not wasm)
    set(CMAKE_SYSTEM_PROCESSOR "wasm32" CACHE STRING "" FORCE)
    # Enable SIMD128 globally for all WASM compilation
    add_compile_options(-msimd128)
    add_link_options(-msimd128)
endif()

# Enable Fortran if requested
if(MLIPCPP_BUILD_FORTRAN)
    enable_language(Fortran)
endif()

# Enable Python bindings via scikit-build
if(SKBUILD)
    set(MLIPCPP_BUILD_PYTHON ON)
endif()

if(MLIPCPP_BUILD_PYTHON)
    find_package(Python 3.10 REQUIRED COMPONENTS Interpreter Development.Module
        OPTIONAL_COMPONENTS Development.SABIModule)
    message(STATUS "Python bindings enabled, setting CMAKE_POSITION_INDEPENDENT_CODE")
    set(CMAKE_POSITION_INDEPENDENT_CODE ON)
endif()

# =============================================================================
# Dependencies (via CPM)
# =============================================================================
set(CPM_DOWNLOAD_VERSION 0.40.2)
if(CPM_SOURCE_CACHE)
    set(CPM_DOWNLOAD_LOCATION "${CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
elseif(DEFINED ENV{CPM_SOURCE_CACHE})
    set(CPM_DOWNLOAD_LOCATION "$ENV{CPM_SOURCE_CACHE}/cpm/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
else()
    set(CPM_DOWNLOAD_LOCATION "${CMAKE_BINARY_DIR}/cmake/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
endif()

if(NOT (EXISTS ${CPM_DOWNLOAD_LOCATION}))
    message(STATUS "Downloading CPM.cmake to ${CPM_DOWNLOAD_LOCATION}")
    file(DOWNLOAD
        https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake
        ${CPM_DOWNLOAD_LOCATION}
        STATUS _cpm_download_status
        TLS_VERIFY ON
    )
    list(GET _cpm_download_status 0 _cpm_download_code)
    if(NOT _cpm_download_code EQUAL 0)
        message(FATAL_ERROR
            "Failed to download CPM.cmake from "
            "https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake "
            "(status: ${_cpm_download_status}). Check your network/proxy, then "
            "delete ${CPM_DOWNLOAD_LOCATION} and re-run cmake.")
    endif()
endif()

include(${CPM_DOWNLOAD_LOCATION})

if(NOT COMMAND CPMAddPackage)
    message(FATAL_ERROR
        "CPM.cmake at ${CPM_DOWNLOAD_LOCATION} is invalid (CPMAddPackage not defined). "
        "Delete that file and re-run cmake to redownload.")
endif()

# ggml (forked version with gradient support)
# Use MLIPCPP_GGML_SOURCE_DIR to point to a local ggml checkout for development
set(MLIPCPP_GGML_SOURCE_DIR "" CACHE PATH "Path to local ggml source directory (for development)")

set(BUILD_SHARED_LIBS OFF CACHE BOOL "" FORCE)
set(GGML_BUILD_TESTS OFF CACHE BOOL "" FORCE)
set(GGML_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)

# Configure ggml backend options based on mlipcpp options
if(MLIPCPP_USE_CUDA)
    set(GGML_CUDA ON CACHE BOOL "" FORCE)
endif()
if(MLIPCPP_USE_HIP)
    set(GGML_HIP ON CACHE BOOL "" FORCE)
endif()
if(MLIPCPP_USE_METAL)
    set(GGML_METAL ON CACHE BOOL "" FORCE)
    set(GGML_METAL_EMBED_LIBRARY ON CACHE BOOL "" FORCE)
else()
    set(GGML_METAL_EMBED_LIBRARY ON CACHE BOOL "" FORCE)
endif()
if(MLIPCPP_USE_VULKAN)
    set(GGML_VULKAN ON CACHE BOOL "" FORCE)
endif()
if(MLIPCPP_USE_WEBGPU)
    set(GGML_WEBGPU ON CACHE BOOL "" FORCE)
    if(EMSCRIPTEN)
        # Default to JSPI; opt in to ASYNCIFY via MLIPCPP_WASM_ASYNCIFY=ON.
        if(MLIPCPP_WASM_ASYNCIFY)
            set(GGML_WEBGPU_JSPI OFF CACHE BOOL "" FORCE)
        else()
            set(GGML_WEBGPU_JSPI ON CACHE BOOL "" FORCE)
        endif()
    endif()
endif()
if(MLIPCPP_USE_SYCL)
    set(GGML_SYCL ON CACHE BOOL "" FORCE)
endif()
if(MLIPCPP_USE_CANN)
    set(GGML_CANN ON CACHE BOOL "" FORCE)
endif()
if(MLIPCPP_USE_BLAS)
    set(GGML_BLAS ON CACHE BOOL "" FORCE)
endif()

# Default to a portable CPU baseline (no -march=native). Wheels and CI need
# this — without it ggml bakes the build host's CPU features into the binary
# and the resulting library SIGILLs on machines with a different (or older)
# CPU. Users wanting host-tuned builds can pass -DGGML_NATIVE=ON explicitly.
if(NOT DEFINED CACHE{GGML_NATIVE})
    set(GGML_NATIVE OFF CACHE BOOL "Build with -march=native (non-portable)")
endif()

# WASM-specific ggml configuration
if(EMSCRIPTEN)
    # SINGLE_FILE=ON embeds the .wasm as base64 inside the .js, which blocks
    # streaming compilation (WebAssembly.instantiateStreaming). Firefox needs
    # streaming to tier up to its optimizing compiler — without it, compute
    # stays stuck in baseline and is ~30x slower than Chrome. Shipping as a
    # separate .wasm fixes this.
    set(GGML_WASM_SINGLE_FILE OFF CACHE BOOL "" FORCE)
    set(GGML_NATIVE OFF CACHE BOOL "" FORCE)
    set(GGML_METAL OFF CACHE BOOL "" FORCE)
    set(GGML_CUDA OFF CACHE BOOL "" FORCE)
    set(GGML_VULKAN OFF CACHE BOOL "" FORCE)
    set(GGML_BLAS OFF CACHE BOOL "" FORCE)
endif()

if(MLIPCPP_GGML_SOURCE_DIR)
    message(STATUS "Using local ggml from: ${MLIPCPP_GGML_SOURCE_DIR}")
    add_subdirectory(${MLIPCPP_GGML_SOURCE_DIR} ${CMAKE_BINARY_DIR}/ggml EXCLUDE_FROM_ALL)
else()
    CPMAddPackage(
        NAME ggml
        GITHUB_REPOSITORY peterspackman/ggml
        GIT_TAG 9be7b7f8
        EXCLUDE_FROM_ALL YES
    )
endif()

# fmt (not needed for WASM builds - logging is disabled)
if(MLIPCPP_USE_SYSTEM_FMT)
    find_package(fmt REQUIRED)
elseif(NOT EMSCRIPTEN)
    CPMAddPackage(
        NAME fmt
        GITHUB_REPOSITORY fmtlib/fmt
        GIT_TAG 11.0.2
        OPTIONS
            "FMT_INSTALL OFF"
        EXCLUDE_FROM_ALL YES
    )
endif()

# nlohmann_json (for graph inference metadata)
CPMAddPackage(
    NAME nlohmann_json
    GITHUB_REPOSITORY nlohmann/json
    VERSION 3.11.3
    OPTIONS
        "JSON_BuildTests OFF"
    EXCLUDE_FROM_ALL YES
)

# =============================================================================
# Library
# =============================================================================
add_subdirectory(src)

# Enable CTest before descending into examples/, so add_test() calls there
# (for c_api_test / cpp_api_test / fortran_api_test) actually register
# rather than being silently dropped.
if(MLIPCPP_BUILD_TESTS)
    enable_testing()
endif()

# =============================================================================
# Examples
# =============================================================================
if(MLIPCPP_BUILD_EXAMPLES)
    add_subdirectory(examples)
endif()

# =============================================================================
# Tests
# =============================================================================
if(MLIPCPP_BUILD_TESTS)
    CPMAddPackage(
        NAME Catch2
        GITHUB_REPOSITORY catchorg/Catch2
        VERSION 3.7.1
        OPTIONS
            "CATCH_BUILD_TESTING OFF"
            "CATCH_INSTALL_DOCS OFF"
            "CATCH_INSTALL_EXTRAS OFF"
        EXCLUDE_FROM_ALL YES
    )

    add_subdirectory(tests)
endif()

# =============================================================================
# Python Bindings
# =============================================================================
if(MLIPCPP_BUILD_PYTHON)
    CPMAddPackage(
        NAME nanobind
        GITHUB_REPOSITORY wjakob/nanobind
        GIT_TAG v2.4.0
        EXCLUDE_FROM_ALL YES
    )

    find_package(nanobind CONFIG REQUIRED)
    include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/PythonBindings.cmake)

    nanobind_add_module(_mlipcpp
        NB_STATIC
        LTO
        STABLE_ABI
        ${CMAKE_CURRENT_SOURCE_DIR}/src/api/python/mlipcpp_bindings.cpp
    )

    target_link_libraries(_mlipcpp PRIVATE mlipcpp fmt::fmt)
    target_include_directories(_mlipcpp PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
    target_compile_definitions(_mlipcpp PRIVATE VERSION_INFO=${PROJECT_VERSION})

    install(TARGETS _mlipcpp LIBRARY DESTINATION mlipcpp)
    install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/python/mlipcpp/
        DESTINATION mlipcpp
        FILES_MATCHING PATTERN "*.py"
    )

    add_nanobind_stubs(_mlipcpp _mlipcpp ${CMAKE_CURRENT_BINARY_DIR}/stubs)
endif()

# =============================================================================
# Installation
# =============================================================================
if(MLIPCPP_INSTALL)
    include(GNUInstallDirs)
    include(CMakePackageConfigHelpers)

    # Install headers
    install(DIRECTORY include/
        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
        FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp"
    )

    # Targets to export
    set(MLIPCPP_EXPORT_TARGETS mlipcpp)
    if(MLIPCPP_BUILD_FORTRAN)
        list(APPEND MLIPCPP_EXPORT_TARGETS mlipcpp_fortran)

        # Install Fortran module files
        install(DIRECTORY ${CMAKE_BINARY_DIR}/fortran_modules/
            DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/mlipcpp/fortran
            FILES_MATCHING PATTERN "*.mod"
        )
    endif()

    # Install library targets
    install(TARGETS ${MLIPCPP_EXPORT_TARGETS}
        EXPORT mlipcppTargets
        LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
        ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
        RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
        INCLUDES DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
    )

    # Export targets
    install(EXPORT mlipcppTargets
        FILE mlipcppTargets.cmake
        NAMESPACE mlipcpp::
        DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/mlipcpp
    )

    # Package config file
    configure_package_config_file(
        ${CMAKE_CURRENT_SOURCE_DIR}/cmake/mlipcppConfig.cmake.in
        ${CMAKE_CURRENT_BINARY_DIR}/mlipcppConfig.cmake
        INSTALL_DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/mlipcpp
    )

    # Version file
    write_basic_package_version_file(
        ${CMAKE_CURRENT_BINARY_DIR}/mlipcppConfigVersion.cmake
        VERSION ${PROJECT_VERSION}
        COMPATIBILITY SameMajorVersion
    )

    # Install config files
    install(FILES
        ${CMAKE_CURRENT_BINARY_DIR}/mlipcppConfig.cmake
        ${CMAKE_CURRENT_BINARY_DIR}/mlipcppConfigVersion.cmake
        DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/mlipcpp
    )

    # Export from build tree (for use without installing)
    export(EXPORT mlipcppTargets
        FILE ${CMAKE_CURRENT_BINARY_DIR}/mlipcppTargets.cmake
        NAMESPACE mlipcpp::
    )
endif()

# =============================================================================
# WASM Build
# =============================================================================
if(EMSCRIPTEN)
    # Create WASM module with embind bindings
    add_executable(mlipcpp_wasm
        ${CMAKE_CURRENT_SOURCE_DIR}/src/api/wasm/mlipcpp_wasm.cpp
    )

    target_link_libraries(mlipcpp_wasm PRIVATE mlipcpp)
    target_include_directories(mlipcpp_wasm PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)

    set(MLIPCPP_WASM_LINK_FLAGS
        "-s WASM=1 "
        "-s MODULARIZE=1 "
        "-s EXPORT_ES6=1 "
        "-s EXPORT_NAME='createMlipcpp' "
        "-s ALLOW_MEMORY_GROWTH=1 "
        "-s MAXIMUM_MEMORY=4GB "
        "-s STACK_SIZE=1MB "
        "-s EXPORTED_RUNTIME_METHODS=['FS','cwrap','ccall'] "
        "-s FORCE_FILESYSTEM=1 "
        "--bind "
        "-O3 "
    )

    # When WebGPU is enabled, ggml-webgpu already propagates -sJSPI or
    # -sASYNCIFY via INTERFACE link options. Add a larger async stack for
    # ASYNCIFY; JSPI doesn't need it.
    if(MLIPCPP_USE_WEBGPU AND MLIPCPP_WASM_ASYNCIFY)
        # NOTE: ASYNCIFY_IGNORE_INDIRECT=1 would narrow the transform nicely
        # (ggml's backend iface dispatch is indirect), but it breaks at
        # runtime with "indirect call to null" — some indirect call reaches
        # a function that legitimately needs to suspend (likely inside
        # emdawnwebgpu's callback plumbing). Leave it off until we can
        # enumerate those with ASYNCIFY_ADD.
        list(APPEND MLIPCPP_WASM_LINK_FLAGS "-s ASYNCIFY_STACK_SIZE=65536 ")
    endif()

    string(REPLACE ";" "" MLIPCPP_WASM_LINK_FLAGS "${MLIPCPP_WASM_LINK_FLAGS}")

    set_target_properties(mlipcpp_wasm PROPERTIES
        SUFFIX ".js"
        RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin
        LINK_FLAGS "${MLIPCPP_WASM_LINK_FLAGS}"
    )

    # Install WASM output
    install(FILES
        ${CMAKE_BINARY_DIR}/bin/mlipcpp_wasm.js
        DESTINATION ${CMAKE_INSTALL_DATADIR}/mlipcpp/wasm
    )
endif()
