cmake_minimum_required(VERSION 3.25)
project(sovereign-engine LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

# ─── Options ──────────────────────────────────────────────────────────────────
option(SOVEREIGN_BUILD_TESTS "Build tests" ON)
option(SOVEREIGN_BUILD_PYTHON "Build Python bindings" ON)
option(SOVEREIGN_ENABLE_AVX512 "Enable AVX-512 fallback kernels" ON)
option(SOVEREIGN_ENABLE_ASAN "Enable AddressSanitizer" OFF)

# ─── Compiler Flags & AddressSanitizer ────────────────────────────────────────
if(MSVC)
    add_compile_options(/W4)
    if(SOVEREIGN_ENABLE_ASAN)
        add_compile_options(/fsanitize=address)
        add_link_options(/fsanitize=address)
    endif()
else()
    add_compile_options(-Wall -Wextra -Wno-unused-parameter)
    if(SOVEREIGN_ENABLE_ASAN)
        add_compile_options(-fsanitize=address -fno-omit-frame-pointer)
        add_link_options(-fsanitize=address)
    endif()
endif()

# ─── Global Compile Definitions ───────────────────────────────────────────────
add_compile_definitions(
    SOVEREIGN_VERSION_MAJOR=0
    SOVEREIGN_VERSION_MINOR=1
    SOVEREIGN_VERSION_PATCH=0
)

if(WIN32)
    add_compile_definitions(WIN32_LEAN_AND_MEAN NOMINMAX)
endif()

if(SOVEREIGN_ENABLE_AVX512)
    add_compile_definitions(SOVEREIGN_ENABLE_AVX512=1)
endif()

# ─── Dependencies ─────────────────────────────────────────────────────────────
find_package(Vulkan)

include(FetchContent)

if(NOT Vulkan_FOUND)
    message(STATUS "Vulkan SDK not found. Setting up SDK-free build with Vulkan-Headers...")
    FetchContent_Declare(
        vulkan_headers
        GIT_REPOSITORY https://github.com/KhronosGroup/Vulkan-Headers.git
        GIT_TAG        v1.3.283
        GIT_SHALLOW    ON
    )
    FetchContent_MakeAvailable(vulkan_headers)
    
    add_library(Vulkan::Vulkan INTERFACE IMPORTED)
    target_link_libraries(Vulkan::Vulkan INTERFACE Vulkan::Headers)
    if(UNIX AND NOT APPLE)
        target_link_libraries(Vulkan::Vulkan INTERFACE ${CMAKE_DL_LIBS})
    endif()
endif()

if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/vk_mem_alloc.h")
    message(STATUS "vk_mem_alloc.h not found. Downloading from GPUOpen GitHub...")
    file(DOWNLOAD
        "https://raw.githubusercontent.com/GPUOpen-LibrariesAndSDKs/VulkanMemoryAllocator/master/include/vk_mem_alloc.h"
        "${CMAKE_CURRENT_SOURCE_DIR}/third_party/vk_mem_alloc.h"
        SHOW_PROGRESS
        STATUS download_status
    )
    list(GET download_status 0 status_code)
    if(NOT status_code EQUAL 0)
        message(FATAL_ERROR "Failed to download vk_mem_alloc.h: ${download_status}")
    endif()
endif()

# ─── Shader Compilation (GLSL -> SPIR-V) ──────────────────────────────────────
find_program(GLSLC_EXECUTABLE glslc)

set(SHADER_SPV_DIR "${CMAKE_BINARY_DIR}/shaders")
file(MAKE_DIRECTORY "${SHADER_SPV_DIR}")

set(SHADER_SOURCES
    shaders/rmsnorm.comp
    shaders/matmul_int4.comp
    shaders/attention_gqa.comp
    shaders/silu_gate.comp
    shaders/sampler.comp
)

if(GLSLC_EXECUTABLE)
    message(STATUS "glslc compiler found at: ${GLSLC_EXECUTABLE}. Recompiling GLSL shaders...")
    foreach(SHADER_SRC ${SHADER_SOURCES})
        get_filename_component(SHADER_NAME ${SHADER_SRC} NAME_WE)
        set(SPV_FILE "${SHADER_SPV_DIR}/${SHADER_NAME}.spv")
        add_custom_command(
            OUTPUT ${SPV_FILE}
            COMMAND ${GLSLC_EXECUTABLE} -O ${CMAKE_CURRENT_SOURCE_DIR}/${SHADER_SRC} -o ${SPV_FILE}
            DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/${SHADER_SRC}
            COMMENT "Compiling shader ${SHADER_SRC} to SPIR-V"
        )
        list(APPEND SHADER_SPV_FILES ${SPV_FILE})
    endforeach()
else()
    message(STATUS "glslc compiler not found. Using precompiled SPIR-V shaders from shaders/compiled/ ...")
    foreach(SHADER_SRC ${SHADER_SOURCES})
        get_filename_component(SHADER_NAME ${SHADER_SRC} NAME_WE)
        set(SPV_FILE "${SHADER_SPV_DIR}/${SHADER_NAME}.spv")
        add_custom_command(
            OUTPUT ${SPV_FILE}
            COMMAND ${CMAKE_COMMAND} -E copy
                    ${CMAKE_CURRENT_SOURCE_DIR}/shaders/compiled/${SHADER_NAME}.spv
                    ${SPV_FILE}
            DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/shaders/compiled/${SHADER_NAME}.spv
            COMMENT "Copying precompiled shader ${SHADER_NAME}.spv"
        )
        list(APPEND SHADER_SPV_FILES ${SPV_FILE})
    endforeach()
endif()

add_custom_target(compile_shaders DEPENDS ${SHADER_SPV_FILES})

# ─── Sovereign Core Library ───────────────────────────────────────────────────
add_library(sovereign_core STATIC
    src/vulkan/vulkan_context.cpp
    src/format/format.cpp
    src/compute/kv_cache.cpp
    src/inference/engine.cpp
    src/quantizer/quantizer.cpp
    src/memory/memory_manager.cpp
    third_party/volk/volk.c
)

target_include_directories(sovereign_core PUBLIC
    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
    $<INSTALL_INTERFACE:include>
    ${CMAKE_CURRENT_SOURCE_DIR}/third_party
)

target_link_libraries(sovereign_core PUBLIC Vulkan::Vulkan)

# Add dependencies for shader outputs
add_dependencies(sovereign_core compile_shaders)

# Define shader directory path macro (normalised for Windows/C++ path string compatibility)
string(REPLACE "\\" "/" SHADER_SPV_DIR_NORM "${SHADER_SPV_DIR}")
target_compile_definitions(sovereign_core PRIVATE SOVEREIGN_SHADER_DIR="${SHADER_SPV_DIR_NORM}")

# ─── Converter CLI Tool ───────────────────────────────────────────────────────
add_executable(sovereign-convert
    tools/converter/main.cpp
)
target_link_libraries(sovereign-convert PRIVATE sovereign_core)

# ─── Python Bindings ──────────────────────────────────────────────────────────
if(SOVEREIGN_BUILD_PYTHON)
    FetchContent_Declare(
        pybind11
        GIT_REPOSITORY https://github.com/pybind/pybind11.git
        GIT_TAG        v2.11.1
        GIT_SHALLOW    ON
    )
    FetchContent_MakeAvailable(pybind11)

    pybind11_add_module(sovereign_inference bindings/python/sovereign_py.cpp)
    target_link_libraries(sovereign_inference PRIVATE sovereign_core)
    install(TARGETS sovereign_inference DESTINATION .)
endif()

# ─── Unit Tests ───────────────────────────────────────────────────────────────
if(SOVEREIGN_BUILD_TESTS)
    FetchContent_Declare(
        doctest
        GIT_REPOSITORY https://github.com/doctest/doctest.git
        GIT_TAG        v2.4.11
        GIT_SHALLOW    ON
    )
    FetchContent_MakeAvailable(doctest)

    enable_testing()
    add_subdirectory(tests)
endif()
