# Builds the gk-diffuser engine (the diffusion runtime + CLI running on gk)
# from the vendored source tree as a single static binary and installs it into
# the gk_diffuser package, where the Python server spawns it per generation.
# Modeled on the ggk and gguf-diffuser CMakeLists.txt; this is the diffusion
# half of the ggk split.
cmake_minimum_required(VERSION 3.15)
project(gk_diffuser C CXX)

option(GK_DIFFUSER_BUILD "Build the gk-diffuser engine and install it alongside the python package" ON)

if(NOT GK_DIFFUSER_BUILD)
    return()
endif()

if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
    set(CMAKE_BUILD_TYPE Release CACHE STRING "Build type" FORCE)
endif()

# ── Locate the engine source tree ────────────────────────────────────────────
set(GK_DIFFUSER_ENGINE_DIR "" CACHE PATH "Path to the gk-diffuser engine source tree")
if(NOT GK_DIFFUSER_ENGINE_DIR AND DEFINED ENV{GK_DIFFUSER_ENGINE_DIR})
    set(GK_DIFFUSER_ENGINE_DIR "$ENV{GK_DIFFUSER_ENGINE_DIR}")
endif()
if(NOT GK_DIFFUSER_ENGINE_DIR AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/vendor/diffusion/CMakeLists.txt")
    set(GK_DIFFUSER_ENGINE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/vendor/diffusion")
endif()
if(NOT GK_DIFFUSER_ENGINE_DIR OR NOT EXISTS "${GK_DIFFUSER_ENGINE_DIR}/diffusion/CMakeLists.txt")
    message(FATAL_ERROR
        "gk-diffuser engine source not found. Provide it via one of:\n"
        "  - the vendored copy at vendor/diffusion (see scripts/vendor_engine.py)\n"
        "  - -DGK_DIFFUSER_ENGINE_DIR=/path/to/engine or the same env var")
endif()
get_filename_component(GK_DIFFUSER_ENGINE_DIR "${GK_DIFFUSER_ENGINE_DIR}" ABSOLUTE)
message(STATUS "gk-diffuser: engine source at ${GK_DIFFUSER_ENGINE_DIR}")

# The engine picks its sources up with file(GLOB "src/*.cpp"), which also
# matches macOS AppleDouble sidecars ("._foo.cpp" — binary metadata, not
# source). Left in place they are handed to the compiler and the build dies
# hundreds of files in, with an error that says nothing about where they came
# from. Catch them at configure time instead.
file(GLOB_RECURSE _gk_diffuser_appledouble
     LIST_DIRECTORIES false
     "${GK_DIFFUSER_ENGINE_DIR}/._*" "${GK_DIFFUSER_ENGINE_DIR}/*/._*")
if(_gk_diffuser_appledouble)
    list(LENGTH _gk_diffuser_appledouble _gk_diffuser_appledouble_count)
    list(GET _gk_diffuser_appledouble 0 _gk_diffuser_appledouble_first)
    message(FATAL_ERROR
        "the engine tree carries ${_gk_diffuser_appledouble_count} macOS AppleDouble file(s), "
        "e.g. ${_gk_diffuser_appledouble_first}\n"
        "They match the engine's source globs and will fail the compile. Remove them:\n"
        "  find \"${GK_DIFFUSER_ENGINE_DIR}\" \\( -name '._*' -o -name '.DS_Store' \\) -delete")
endif()

# ── Option helper ────────────────────────────────────────────────────────────
# Values are read from the environment as well as -D, because `pip install`
# makes environment variables far easier to pass than CMake defines:
#   GK_DIFFUSER_CUDA=1 pip install .
#   CMAKE_ARGS="-DGK_DIFFUSER_CUDA=ON" pip install .
function(gk_diffuser_option name docstring default)
    if(DEFINED ENV{${name}} AND NOT "$ENV{${name}}" STREQUAL "")
        set(default "$ENV{${name}}")
    endif()
    option(${name} "${docstring}" "${default}")
endfunction()

# ── Windows / MinGW toolchain fixes ──────────────────────────────────────────
if(WIN32 AND NOT MSVC)
    # MinGW-w64's headers pin _WIN32_WINNT to a Windows 7-era value before the
    # engine gets a chance to set it. Set it up front, matching what the
    # engine expects everywhere else.
    add_compile_definitions(_WIN32_WINNT=0x0A00 WINVER=0x0A00 NTDDI_VERSION=0x0A000000)

    # A MinGW-built binary otherwise needs libstdc++-6.dll, libgcc_s_seh-1.dll
    # and libwinpthread-1.dll from the toolchain at run time. The wheel ships a
    # single file and Python spawns it without MSYS2 on PATH, so link the
    # runtime in.
    gk_diffuser_option(GK_DIFFUSER_MINGW_STATIC
        "Statically link the MinGW runtime into the engine binary" ON)
    if(GK_DIFFUSER_MINGW_STATIC)
        # gk runs its own thread pool rather than OpenMP, so -static has no
        # libgomp import library to fight here.
        add_link_options(-static-libgcc -static-libstdc++ -static)
    endif()
endif()

# ── Shape the engine build for a single-file wheel ───────────────────────────
# Static libs (gk, ggml-compat, diffusion) linked straight into the executable
# keep the wheel to one binary with no rpath/DLL handling.
set(BUILD_SHARED_LIBS OFF CACHE BOOL "gk-diffuser: static engine for a single-file install" FORCE)
set(SD_BUILD_SHARED_LIBS OFF CACHE BOOL "gk-diffuser: static libdiffusion" FORCE)
set(SD_BUILD_CLI ON CACHE BOOL "gk-diffuser: the CLI is the bundled artifact" FORCE)

# ── GPU / accelerator backends ───────────────────────────────────────────────
# One switch per backend drives the whole engine; the gk kernels evaluate every
# graph the runtime builds. Anything else gk supports can still be passed
# straight through as -DGK_*.
#
# A CUDA build works its own architecture list out from nvcc and the installed
# GPUs, and embeds PTX for the newest, so an unlisted card JITs rather than
# failing. A wheel built on one machine for another should still say what it is
# targeting: -DGK_CUDA_ARCHITECTURES="75;86;89;120".
gk_diffuser_option(GK_DIFFUSER_CUDA   "Build the engine with the gk CUDA backend"   OFF)
gk_diffuser_option(GK_DIFFUSER_HIP    "Build the engine with the gk HIP backend"    OFF)
gk_diffuser_option(GK_DIFFUSER_VULKAN "Build the engine with the gk Vulkan backend" OFF)
if(APPLE)  # matches the engine's own default; APPLE is "" elsewhere, not a valid default
    set(_gk_metal_default ON)
else()
    set(_gk_metal_default OFF)
endif()
gk_diffuser_option(GK_DIFFUSER_METAL  "Build the engine with the gk Metal backend"  ${_gk_metal_default})

set(GK_CUDA   ${GK_DIFFUSER_CUDA}   CACHE BOOL "" FORCE)
set(GK_HIP    ${GK_DIFFUSER_HIP}    CACHE BOOL "" FORCE)
set(GK_VULKAN ${GK_DIFFUSER_VULKAN} CACHE BOOL "" FORCE)
set(GK_METAL  ${GK_DIFFUSER_METAL}  CACHE BOOL "" FORCE)

set(GK_DIFFUSER_ENABLED_BACKENDS "")
foreach(backend IN ITEMS CUDA HIP VULKAN METAL)
    if(GK_DIFFUSER_${backend})
        list(APPEND GK_DIFFUSER_ENABLED_BACKENDS ${backend})
    endif()
endforeach()
if(GK_DIFFUSER_ENABLED_BACKENDS)
    string(REPLACE ";" ", " _backend_list "${GK_DIFFUSER_ENABLED_BACKENDS}")
    message(STATUS "gk-diffuser: accelerator backends requested: ${_backend_list}")
else()
    message(STATUS "gk-diffuser: CPU-only build (enable e.g. -DGK_DIFFUSER_CUDA=ON for GPU)")
endif()

# EXCLUDE_FROM_ALL keeps every target the engine defines out of the default
# build — the dependency below pulls in the CLI binary and just the libraries
# it links against.
add_subdirectory("${GK_DIFFUSER_ENGINE_DIR}" "${CMAKE_BINARY_DIR}/gk-diffuser-engine" EXCLUDE_FROM_ALL)

if(NOT TARGET diffusion-cli)
    message(FATAL_ERROR "engine did not define the diffusion-cli target — is ${GK_DIFFUSER_ENGINE_DIR} the gk-diffuser engine source tree?")
endif()

add_custom_target(gk_diffuser_engine ALL)
add_dependencies(gk_diffuser_engine diffusion-cli)

# Install into the wheel (SKBUILD_PLATLIB_DIR) and also into the source tree
# so editable installs (`pip install -e .`) find the binary — the same dual
# install ggk and the gguf-* packages use.
install(
    PROGRAMS $<TARGET_FILE:diffusion-cli>
    DESTINATION "${CMAKE_CURRENT_SOURCE_DIR}/src/gk_diffuser/bin"
)
if(DEFINED SKBUILD_PLATLIB_DIR)
    install(
        PROGRAMS $<TARGET_FILE:diffusion-cli>
        DESTINATION "${SKBUILD_PLATLIB_DIR}/gk_diffuser/bin"
    )
endif()
