cmake_minimum_required(VERSION 3.18)
project(openllava VERSION 3.0.0 LANGUAGES C CXX CUDA)

option(OPENLLAVA_BUILD_PYTHON "Build Python bindings" OFF)

set(CMAKE_C_STANDARD 17)
set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CUDA_STANDARD 17)
set(CMAKE_CUDA_STANDARD_REQUIRED ON)

set(CMAKE_POSITION_INDEPENDENT_CODE ON)

# Find packages
find_package(Python3 REQUIRED COMPONENTS Interpreter Development)
find_package(pybind11 REQUIRED)

# CUDA
find_package(CUDAToolkit REQUIRED)
set(CMAKE_CUDA_ARCHITECTURES "80-real;86-real;89-real;90-real" CACHE STRING "CUDA architectures")
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -use_fast_math")
set(CMAKE_CUDA_SEPARABLE_COMPILATION ON)

# OpenLLaVA C++ library (core)
add_library(openllava_core SHARED
    csrc/gpu/streams.cpp
    csrc/gpu/memory_manager.cpp
    csrc/cpu/offload.cpp
    csrc/gpu/projector_kernel.cu
    csrc/gpu/vision_pack.cu
    csrc/gpu/vq_lookup.cu
    csrc/gpu/cross_attention.cu
)

target_include_directories(openllava_core PRIVATE
    ${CMAKE_SOURCE_DIR}/csrc
    ${CMAKE_SOURCE_DIR}/csrc/gpu
    ${CUDAToolkit_INCLUDE_DIRS}
    ${Python3_INCLUDE_DIRS}
    ${pybind11_INCLUDE_DIRS}
)

target_link_libraries(openllava_core PRIVATE
    ${CUDA_LIBRARIES}
    CUDA::cudart
    CUDA::cublas
)

# CPU-only shared library (fallback when no GPU)
add_library(openllava_cpu SHARED
    csrc/cpu/fp32_ops.c
    csrc/cpu/quantize.c
    csrc/cpu/gguf_writer.c
    csrc/cpu/bitnet_cpu.c
)

target_include_directories(openllava_cpu PRIVATE
    ${CMAKE_SOURCE_DIR}/csrc
)

# pybind11 bridge module
pybind11_add_module(_openllava_ext
    csrc/bindings/pybind.cpp
)

target_link_libraries(_openllava_ext PRIVATE
    openllava_core
    openllava_cpu
)

target_include_directories(_openllava_ext PRIVATE
    ${CMAKE_SOURCE_DIR}/csrc
)

# TPU XLA backend (C++ side)
add_library(openllava_tpu SHARED
    csrc/tpu/xla_backend.cpp
)

target_include_directories(openllava_tpu PRIVATE
    ${CMAKE_SOURCE_DIR}/csrc
    ${Python3_INCLUDE_DIRS}
)

# Install targets
include(GNUInstallDirs)
install(TARGETS openllava_core openllava_cpu openllava_tpu
    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
)

# Only install _openllava_ext if we are building Python bindings
if(OPENLLAVA_BUILD_PYTHON)
    install(TARGETS _openllava_ext
        LIBRARY DESTINATION openllava
    )
endif()
