cmake_minimum_required(VERSION 3.21)
project(scene_actor_extraction VERSION 0.1.0 LANGUAGES CXX)

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

# ── Dependencies ──────────────────────────────────────────────────────────────

# KPN++ (pipeline backbone)
set(KPN_BUILD_TESTS    OFF CACHE BOOL "" FORCE)
set(KPN_BUILD_PYTHON   OFF CACHE BOOL "" FORCE)
set(KPN_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
option(SAE_WEB_DEBUG "Enable KPN web debug UI (localhost:9090)" OFF)
if(SAE_WEB_DEBUG)
    set(KPN_WEB_DEBUG ON CACHE BOOL "" FORCE)
endif()
add_subdirectory(external/KPN)

# OpenCV (video decode, image ops, DNN inference, face detection)
find_package(OpenCV 4 REQUIRED COMPONENTS
    core imgproc imgcodecs videoio dnn objdetect highgui)

# ONNX Runtime (SCRFD detector — cv::dnn cannot handle dynamic Shape nodes)
find_library(ORT_LIB onnxruntime REQUIRED
    HINTS /usr/lib /usr/local/lib)
find_path(ORT_INCLUDE onnxruntime_cxx_api.h
    PATH_SUFFIXES onnxruntime
    HINTS /usr/include /usr/local/include
    REQUIRED)
add_library(onnxruntime UNKNOWN IMPORTED)
set_target_properties(onnxruntime PROPERTIES
    IMPORTED_LOCATION "${ORT_LIB}"
    INTERFACE_INCLUDE_DIRECTORIES "${ORT_INCLUDE}")
message(STATUS "ONNX Runtime: ${ORT_LIB}  headers: ${ORT_INCLUDE}")

# TensorRT + CUDA runtime (raw-TRT ArcFace embedder; activated by --arcface-engine).
find_library(NVINFER_LIB nvinfer
    HINTS /usr/lib /usr/local/lib /opt/tensorrt/lib)
find_path(NVINFER_INCLUDE NvInfer.h
    HINTS /usr/include /usr/local/include /opt/tensorrt/include)
find_library(CUDART_LIB cudart
    HINTS /opt/cuda/lib64 /usr/local/cuda/lib64 /usr/lib)
find_path(CUDART_INCLUDE cuda_runtime_api.h
    HINTS /opt/cuda/targets/x86_64-linux/include /opt/cuda/include
          /usr/local/cuda/include /usr/include)
find_library(CUBLAS_LIB cublas
    HINTS /opt/cuda/targets/x86_64-linux/lib /opt/cuda/lib64
          /usr/local/cuda/lib64 /usr/lib)
if(NOT (NVINFER_LIB AND NVINFER_INCLUDE AND CUDART_LIB AND CUDART_INCLUDE AND CUBLAS_LIB))
    message(FATAL_ERROR
        "TensorRT or CUDA runtime not found "
        "(nvinfer=${NVINFER_LIB} headers=${NVINFER_INCLUDE} "
        "cudart=${CUDART_LIB} cublas=${CUBLAS_LIB} headers=${CUDART_INCLUDE})")
endif()
add_library(trt_runtime INTERFACE)
target_include_directories(trt_runtime INTERFACE
    "${NVINFER_INCLUDE}" "${CUDART_INCLUDE}")
target_link_libraries(trt_runtime INTERFACE
    "${NVINFER_LIB}" "${CUDART_LIB}" "${CUBLAS_LIB}")
message(STATUS "TensorRT: ${NVINFER_LIB}  CUDA runtime: ${CUDART_LIB}  cuBLAS: ${CUBLAS_LIB}")

# FFmpeg (NVDEC hardware video decode + swscale colour conversion)
find_package(PkgConfig REQUIRED)
pkg_check_modules(AVFORMAT REQUIRED libavformat)
pkg_check_modules(AVCODEC  REQUIRED libavcodec)
pkg_check_modules(AVUTIL   REQUIRED libavutil)
pkg_check_modules(SWSCALE  REQUIRED libswscale)

add_library(ffmpeg_libs INTERFACE)
target_compile_options(ffmpeg_libs INTERFACE
    ${AVFORMAT_CFLAGS_OTHER} ${AVCODEC_CFLAGS_OTHER}
    ${AVUTIL_CFLAGS_OTHER}   ${SWSCALE_CFLAGS_OTHER})
target_include_directories(ffmpeg_libs INTERFACE
    ${AVFORMAT_INCLUDE_DIRS} ${AVCODEC_INCLUDE_DIRS}
    ${AVUTIL_INCLUDE_DIRS}   ${SWSCALE_INCLUDE_DIRS})
target_link_libraries(ffmpeg_libs INTERFACE
    ${AVFORMAT_LIBRARIES} ${AVCODEC_LIBRARIES}
    ${AVUTIL_LIBRARIES}   ${SWSCALE_LIBRARIES})
message(STATUS "FFmpeg: avformat=${AVFORMAT_VERSION} avcodec=${AVCODEC_VERSION}")

# nlohmann/json (gallery + output serialisation)
include(FetchContent)
FetchContent_Declare(
    nlohmann_json
    GIT_REPOSITORY https://github.com/nlohmann/json.git
    GIT_TAG        v3.11.3
    GIT_SHALLOW    TRUE
)
FetchContent_MakeAvailable(nlohmann_json)

# nanobind (Python bindings for the sae_embed module)
find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
FetchContent_Declare(
    nanobind
    GIT_REPOSITORY https://github.com/wjakob/nanobind.git
    GIT_TAG        v2.4.0
    GIT_SHALLOW    TRUE
)
FetchContent_MakeAvailable(nanobind)

# ── Model paths ───────────────────────────────────────────────────────────────
set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models"
    CACHE PATH "Directory containing ONNX model files")

# ── Shared library: gallery store ─────────────────────────────────────────────
add_library(sae_gallery STATIC
    src/gallery/gallery_store.cpp
    src/gallery/gallery_builder.cpp
)
target_include_directories(sae_gallery PUBLIC src)
target_link_libraries(sae_gallery PUBLIC
    kpn
    ${OpenCV_LIBS}
    nlohmann_json::nlohmann_json
    onnxruntime
    trt_runtime
    ffmpeg_libs
)
target_compile_definitions(sae_gallery PUBLIC
    SAE_MODELS_DIR="${SAE_MODELS_DIR}"
)

# ── embed_faces — image → embedding JSON (used by gallery builder scripts) ────
add_executable(embed_faces src/embed_faces.cpp)
target_link_libraries(embed_faces PRIVATE
    kpn
    ${OpenCV_LIBS}
    nlohmann_json::nlohmann_json
    onnxruntime
    trt_runtime
    ffmpeg_libs
)
target_compile_definitions(embed_faces PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}")

# ── sae_embed — Python module: load SCRFD+ArcFace once, embed many images ───
nanobind_add_module(sae_embed src/python_bindings.cpp)
target_include_directories(sae_embed PRIVATE src)
target_link_libraries(sae_embed PRIVATE
    ${OpenCV_LIBS}
    onnxruntime
)
target_compile_definitions(sae_embed PRIVATE SAE_MODELS_DIR="${SAE_MODELS_DIR}")

# ── analyze — main analysis binary ───────────────────────────────────────────
add_executable(scene_analyze src/main.cpp)
target_link_libraries(scene_analyze PRIVATE sae_gallery)

# ── analyze_debug — same binary with debug frame/crop output ─────────────────
add_executable(scene_analyze_debug src/main.cpp)
target_link_libraries(scene_analyze_debug PRIVATE sae_gallery)
target_compile_definitions(scene_analyze_debug PRIVATE SAE_DEBUG=1)

# ── scene_preview — live annotated display while analysing ───────────────────
add_executable(scene_preview src/scene_preview.cpp)
target_link_libraries(scene_preview PRIVATE sae_gallery)

# ── build_gallery — offline gallery construction tool ────────────────────────
add_executable(build_gallery src/build_gallery.cpp)
target_link_libraries(build_gallery PRIVATE sae_gallery)

# ── Optional: web debug UI for pipeline introspection ────────────────────────
if(SAE_WEB_DEBUG)
    kpn_target_enable_web_debug(scene_analyze)
    kpn_target_enable_web_debug(scene_analyze_debug)
    kpn_target_enable_web_debug(scene_preview)
endif()

message(STATUS "OpenCV ${OpenCV_VERSION} found")
message(STATUS "Models dir: ${SAE_MODELS_DIR}")
