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)
# Accept 4 or 5: the APIs used here are stable across both, and distros have
# begun shipping 5.x as the default (Arch/CachyOS). find_package's version
# argument is a minimum, but OpenCV's config rejects a 5.x install when 4 is
# requested, so probe for 5 first and fall back to 4.
find_package(OpenCV 5 QUIET COMPONENTS
    core imgproc imgcodecs videoio dnn objdetect highgui)
if(NOT OpenCV_FOUND)
    find_package(OpenCV 4 REQUIRED COMPONENTS
        core imgproc imgcodecs videoio dnn objdetect highgui)
endif()
message(STATUS "OpenCV: ${OpenCV_VERSION}")

# ── Model paths ───────────────────────────────────────────────────────────────
# Defined early so the backend object libraries below can embed it.
set(SAE_MODELS_DIR "${CMAKE_SOURCE_DIR}/models"
    CACHE PATH "Directory containing ONNX model files")

# ── Backend selection ─────────────────────────────────────────────────────────
# Two independent compile-time axes. The core application is agnostic to both:
# only the matching backend .cpp (in src/backends/) is compiled, and the backend
# headers (onnxruntime / NvInfer.h / cublas / rocblas) never reach core TUs.
#
#   SAE_INFERENCE_BACKEND  ORT  → SCRFD + ArcFace via ONNX Runtime (.onnx models)
#                          TRT  → SCRFD + ArcFace via raw TensorRT (.engine files)
#   SAE_GEMM_BACKEND       ROCM → gallery similarity GEMM via rocBLAS / HIP
#                          CUDA → gallery similarity GEMM via cuBLAS / CUDA
#                          CPU  → portable reference GEMM (no GPU; CI / testing)
set(SAE_INFERENCE_BACKEND "ORT"  CACHE STRING "Inference backend: ORT | TRT")
set(SAE_GEMM_BACKEND       "ROCM" CACHE STRING "Gallery GEMM backend: ROCM | CUDA | CPU")
set_property(CACHE SAE_INFERENCE_BACKEND PROPERTY STRINGS ORT TRT)
set_property(CACHE SAE_GEMM_BACKEND       PROPERTY STRINGS ROCM CUDA CPU)

# Enable the ORT TensorRT/CUDA execution providers inside the ORT inference
# backend (only meaningful when ORT was built with the TensorRT EP). Off by
# default so ROCm/CPU builds don't reference unavailable EPs.
option(SAE_ORT_TRT_EP "ORT backend: enable TensorRT/CUDA execution providers" OFF)

# Back-compat: a legacy -DSAE_WITH_TRT=ON/OFF seeds the new vars (ON⇒TRT+CUDA,
# OFF⇒ORT+ROCM) unless the user set them explicitly.
if(DEFINED SAE_WITH_TRT)
    if(SAE_WITH_TRT)
        set(SAE_INFERENCE_BACKEND "TRT"  CACHE STRING "" FORCE)
        set(SAE_GEMM_BACKEND      "CUDA" CACHE STRING "" FORCE)
    else()
        set(SAE_INFERENCE_BACKEND "ORT"  CACHE STRING "" FORCE)
        set(SAE_GEMM_BACKEND      "ROCM" CACHE STRING "" FORCE)
    endif()
    message(STATUS "SAE_WITH_TRT=${SAE_WITH_TRT} (legacy) → "
                   "SAE_INFERENCE_BACKEND=${SAE_INFERENCE_BACKEND} "
                   "SAE_GEMM_BACKEND=${SAE_GEMM_BACKEND}")
endif()

if(NOT SAE_INFERENCE_BACKEND MATCHES "^(ORT|TRT)$")
    message(FATAL_ERROR "SAE_INFERENCE_BACKEND must be ORT or TRT (got '${SAE_INFERENCE_BACKEND}')")
endif()
if(NOT SAE_GEMM_BACKEND MATCHES "^(ROCM|CUDA|CPU)$")
    message(FATAL_ERROR "SAE_GEMM_BACKEND must be ROCM, CUDA or CPU (got '${SAE_GEMM_BACKEND}')")
endif()

# CUDA runtime is needed by both TRT inference and CUDA GEMM — find it once.
function(sae_find_cudart)
    if(TARGET cudart_dep)
        return()
    endif()
    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)
    if(NOT (CUDART_LIB AND CUDART_INCLUDE))
        message(FATAL_ERROR "CUDA runtime not found (cudart=${CUDART_LIB} headers=${CUDART_INCLUDE}).")
    endif()
    add_library(cudart_dep INTERFACE)
    target_include_directories(cudart_dep INTERFACE "${CUDART_INCLUDE}")
    target_link_libraries(cudart_dep INTERFACE "${CUDART_LIB}")
    set_property(GLOBAL PROPERTY sae_cudart_found TRUE)
endfunction()

# ── Inference backend dependency: builds the `inference_backend` object lib ────
if(SAE_INFERENCE_BACKEND STREQUAL "ORT")
    find_library(ORT_LIB onnxruntime REQUIRED
        HINTS /usr/lib64/rocm/lib /usr/lib /usr/local/lib)
    find_path(ORT_INCLUDE onnxruntime_cxx_api.h
        PATH_SUFFIXES onnxruntime
        HINTS /usr/lib64/rocm/include/onnxruntime /usr/include/onnxruntime /usr/local/include/onnxruntime
              /usr/lib64/rocm/include /usr/include /usr/local/include
        REQUIRED)
    # The include directive is <onnxruntime/onnxruntime_cxx_api.h>, so we need the
    # parent of the onnxruntime/ subdirectory on the include path.
    get_filename_component(ORT_INCLUDE_PARENT "${ORT_INCLUDE}" DIRECTORY)
    if(NOT EXISTS "${ORT_INCLUDE_PARENT}/onnxruntime")
        set(ORT_INCLUDE_PARENT "${ORT_INCLUDE}")
    endif()
    message(STATUS "Inference backend: ORT  (${ORT_LIB}  headers: ${ORT_INCLUDE_PARENT})")

    add_library(inference_backend OBJECT src/backends/ort_backend.cpp)
    set_target_properties(inference_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
    target_include_directories(inference_backend PRIVATE src "${ORT_INCLUDE_PARENT}")
    target_link_libraries(inference_backend PRIVATE ${OpenCV_LIBS} "${ORT_LIB}")
    target_compile_definitions(inference_backend PRIVATE
        SAE_MODELS_DIR="${SAE_MODELS_DIR}"
        $<$<BOOL:${SAE_ORT_TRT_EP}>:SAE_ORT_WITH_TRT_EP>)
else()  # TRT
    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)
    if(NOT (NVINFER_LIB AND NVINFER_INCLUDE))
        message(FATAL_ERROR
            "TensorRT not found (nvinfer=${NVINFER_LIB} headers=${NVINFER_INCLUDE}). "
            "Pass -DSAE_INFERENCE_BACKEND=ORT to load .onnx models without TensorRT.")
    endif()
    sae_find_cudart()
    message(STATUS "Inference backend: TRT  (${NVINFER_LIB})")

    add_library(inference_backend OBJECT src/backends/trt_backend.cpp)
    set_target_properties(inference_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
    target_include_directories(inference_backend PRIVATE src "${NVINFER_INCLUDE}")
    target_link_libraries(inference_backend PRIVATE
        ${OpenCV_LIBS} "${NVINFER_LIB}" cudart_dep)
    target_compile_definitions(inference_backend PRIVATE
        SAE_MODELS_DIR="${SAE_MODELS_DIR}")
endif()

# ── GEMM backend dependency: builds the `gemm_backend` object lib ──────────────
if(SAE_GEMM_BACKEND STREQUAL "CPU")
    # Portable reference GEMM: no GPU libraries, no headers. Used for CI and as
    # the correctness oracle for the CUDA/ROCm backends.
    message(STATUS "GEMM backend: CPU  (portable reference, no GPU)")

    add_library(gemm_backend OBJECT src/backends/gemm_backend.cpp)
    set_target_properties(gemm_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
    target_include_directories(gemm_backend PRIVATE src)
    target_compile_definitions(gemm_backend PRIVATE SAE_GEMM_CPU)
elseif(SAE_GEMM_BACKEND STREQUAL "CUDA")
    find_library(CUBLAS_LIB cublas
        HINTS /opt/cuda/targets/x86_64-linux/lib /opt/cuda/lib64
              /usr/local/cuda/lib64 /usr/lib)
    if(NOT CUBLAS_LIB)
        message(FATAL_ERROR "cuBLAS not found (cublas=${CUBLAS_LIB}).")
    endif()
    sae_find_cudart()
    message(STATUS "GEMM backend: CUDA  (${CUBLAS_LIB})")

    add_library(gemm_backend OBJECT src/backends/gemm_backend.cpp)
    set_target_properties(gemm_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
    target_include_directories(gemm_backend PRIVATE src)
    target_link_libraries(gemm_backend PRIVATE "${CUBLAS_LIB}" cudart_dep)
    target_compile_definitions(gemm_backend PRIVATE SAE_GEMM_CUDA)
else()  # ROCM
    find_library(ROCBLAS_LIB rocblas
        HINTS /usr/lib64/rocm/lib /usr/lib64 /usr/local/lib)
    find_path(ROCBLAS_INCLUDE rocblas/rocblas.h
        HINTS /usr/lib64/rocm/include /usr/include /usr/local/include)
    find_library(HIP_LIB amdhip64
        HINTS /usr/lib64/rocm/lib /usr/lib64 /usr/local/lib)
    find_path(HIP_INCLUDE hip/hip_runtime_api.h
        HINTS /usr/lib64/rocm/include /usr/include /usr/local/include)
    if(NOT (ROCBLAS_LIB AND ROCBLAS_INCLUDE AND HIP_LIB AND HIP_INCLUDE))
        message(FATAL_ERROR
            "rocBLAS or HIP runtime not found "
            "(rocblas=${ROCBLAS_LIB} headers=${ROCBLAS_INCLUDE} "
            "hip=${HIP_LIB} headers=${HIP_INCLUDE}). "
            "Install rocblas-devel and hip-devel (or pass -DSAE_GEMM_BACKEND=CUDA).")
    endif()
    message(STATUS "GEMM backend: ROCM  (${ROCBLAS_LIB})")

    add_library(gemm_backend OBJECT src/backends/gemm_backend.cpp)
    set_target_properties(gemm_backend PROPERTIES POSITION_INDEPENDENT_CODE ON)
    target_include_directories(gemm_backend PRIVATE src "${ROCBLAS_INCLUDE}" "${HIP_INCLUDE}")
    target_link_libraries(gemm_backend PRIVATE "${ROCBLAS_LIB}" "${HIP_LIB}")
    # HIP headers require the platform to be declared explicitly when compiled with g++.
    target_compile_definitions(gemm_backend PRIVATE SAE_GEMM_ROCM __HIP_PLATFORM_AMD__)
endif()

# FFmpeg (hwaccel video decode: CUDA/VAAPI, runtime-detected + swscale colour
# conversion). Hwaccel support is built into libavcodec/libavutil; no extra
# libraries are needed here.
# libswresample is the audio side of the same dependency — downmix + resample
# for the audio signature (IR-004, src/audio_signature.cpp). Not a new project
# dependency: it ships with the libav* set already required above.
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)
pkg_check_modules(SWRESAMPLE REQUIRED libswresample)

add_library(ffmpeg_libs INTERFACE)
target_compile_options(ffmpeg_libs INTERFACE
    ${AVFORMAT_CFLAGS_OTHER} ${AVCODEC_CFLAGS_OTHER}
    ${AVUTIL_CFLAGS_OTHER}   ${SWSCALE_CFLAGS_OTHER}
    ${SWRESAMPLE_CFLAGS_OTHER})
target_include_directories(ffmpeg_libs INTERFACE
    ${AVFORMAT_INCLUDE_DIRS} ${AVCODEC_INCLUDE_DIRS}
    ${AVUTIL_INCLUDE_DIRS}   ${SWSCALE_INCLUDE_DIRS}
    ${SWRESAMPLE_INCLUDE_DIRS})
target_link_libraries(ffmpeg_libs INTERFACE
    ${AVFORMAT_LIBRARIES} ${AVCODEC_LIBRARIES}
    ${AVUTIL_LIBRARIES}   ${SWSCALE_LIBRARIES}
    ${SWRESAMPLE_LIBRARIES})
message(STATUS "FFmpeg: avformat=${AVFORMAT_VERSION} avcodec=${AVCODEC_VERSION} "
               "swresample=${SWRESAMPLE_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 + compiled-in backends ──────────────────────
# The backend object libraries carry their own ORT/TRT/CUDA/ROCm linkage and
# headers; sae_gallery re-exports those object files so every binary that links
# sae_gallery gets the chosen backend without ever seeing its headers.
# HDF5 (C++) — gallery fast-load path + embedding dump. Found here so sae_gallery
# (gallery_store.cpp) can link it; scene_analyze/dump_embeddings reuse the same vars.
find_package(HDF5 REQUIRED COMPONENTS CXX)

add_library(sae_gallery STATIC
    src/gallery/gallery_store.cpp
    src/gallery/gallery_builder.cpp
    src/audio_signature.cpp        # IR-004 — content-derived audio signature
    src/gallery/embedder_stamp.cpp # GR-004 — gallery/embedder binding
)
set_target_properties(sae_gallery PROPERTIES POSITION_INDEPENDENT_CODE ON)
target_include_directories(sae_gallery PUBLIC src ${HDF5_INCLUDE_DIRS})
target_link_libraries(sae_gallery PUBLIC
    kpn
    ${OpenCV_LIBS}
    nlohmann_json::nlohmann_json
    inference_backend
    gemm_backend
    ffmpeg_libs
    ${HDF5_CXX_LIBRARIES}
)
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 sae_gallery)

# ── sae_embed — Python module: load SCRFD+ArcFace once, embed many images ───
nanobind_add_module(sae_embed src/python_bindings.cpp)
target_link_libraries(sae_embed PRIVATE sae_gallery)

# ── sae_kpn — Python module: run the real downstream nodes over dumped embeddings ─
# Assembles face_tracker/identity_matcher/scene_tracker in a Python-driven KPN
# network (KPN_BUILD_PYTHON is enabled per-TU inside the .cpp). Powers the
# threshold-sweep optimizer in scripts/optimizer/.
nanobind_add_module(sae_kpn src/kpn_bindings.cpp)
target_link_libraries(sae_kpn PRIVATE sae_gallery)

# HDF5 already found above (before sae_gallery); vars HDF5_CXX_LIBRARIES / _INCLUDE_DIRS
# are reused by scene_analyze / dump_embeddings below.

# ── analyze — main analysis binary ───────────────────────────────────────────
add_executable(scene_analyze src/main.cpp)
target_link_libraries(scene_analyze PRIVATE sae_gallery ${HDF5_CXX_LIBRARIES})
target_include_directories(scene_analyze PRIVATE ${HDF5_INCLUDE_DIRS})

# ── 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 ${HDF5_CXX_LIBRARIES})
target_include_directories(scene_analyze_debug PRIVATE ${HDF5_INCLUDE_DIRS})
target_compile_definitions(scene_analyze_debug PRIVATE SAE_DEBUG=1)

# ── dump_embeddings — standalone embedding dumper, NO gallery/matcher ─────────
# Front-half only (decode→detect→align→embed→HDF5) for the optimizer replay corpus
# and model bake-off. Skips gallery load + calibration (~24s/run faster).
add_executable(dump_embeddings src/dump_embeddings.cpp)
target_link_libraries(dump_embeddings PRIVATE sae_gallery ${HDF5_CXX_LIBRARIES})
target_include_directories(dump_embeddings PRIVATE ${HDF5_INCLUDE_DIRS})

# ── 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()

# ── Tests ─────────────────────────────────────────────────────────────────────
option(SAE_BUILD_TESTS "Build unit tests (GPU-free)" OFF)
if(SAE_BUILD_TESTS)
    enable_testing()
    add_subdirectory(tests)
endif()

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