# This copyright notice applies to this file only
#
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: MIT
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
# DEALINGS IN THE SOFTWARE.

cmake_minimum_required(VERSION 3.21)

project(PyNvVideoCodec VERSION 2.0)

set(PYTHON_MODULE_PREFIX "PyNvVideoCodec")

set(PYNVVIDEOCODEC_VERSION_MAJOR 2)
set(PYNVVIDEOCODEC_VERSION_MINOR 0)

if(POLICY CMP0135)
    cmake_policy(SET CMP0135 NEW)
endif()

set(CMAKE_POLICY_VERSION_MINIMUM 3.5)

configure_file("inc/Version.hpp.in" "pynvcode_version.h")

find_package(Python ${REQUESTED_PYTHON_VERSION} COMPONENTS Interpreter Development.Module)

option(FETCHCONTENT_QUIET OFF)
include(FetchContent)
fetchcontent_declare(
    pybind11
    URL https://github.com/pybind/pybind11/archive/v2.10.0.tar.gz
    URL_HASH SHA256=eacf582fa8f696227988d08cfc46121770823839fe9e301a20fbce67e7cd70ec
)
fetchcontent_makeavailable(pybind11)

if(DEMUX_ONLY)
    set(PY_SOURCES
        src/PyNvDemuxer.cpp
        src/NvDemuxer.cpp
    )
    set(PY_HDRS
        inc
        ../VideoCodecSDKUtils/helper_classes/Utils
    )
else()
    set(PY_SOURCES
        src/DLPackUtils.cpp
        src/ExternalBuffer.cpp
        src/PyNvVideoCodec.cpp
        src/PyNvDemuxer.cpp
        src/PyNvEncoder.cpp
        src/NvDemuxer.cpp
        src/PyCAIMemoryView.cpp
        src/PyNvDecoder.cpp
        src/NvEncoderClInterface.cpp
        src/SimpleDecoder.cpp
        src/PyNvSimpleDecoder.cpp
        src/DecoderCommon.cpp
        src/ThreadedDecoder.cpp
        src/PyNvThreadedDecoder.cpp
        src/PyNvSimpleTranscoder.cpp
        src/SimpleTranscoder.cpp
        src/SeekUtils.cpp
        ../VideoCodecSDKUtils/helper_classes/NvCodec/NvEncoder/NvEncoderCuda.cpp
    )
    set(PY_HDRS
        inc
        ../VideoCodecSDKUtils/helper_classes/NvCodec/NvEncoder
        ${FFMPEG_DIR}/include/
        utils
    )
endif()

# DLPack setup
include(FetchContent)
set(PYNVC_DLPACK_SOURCE_DIR "$ENV{DLPACK_SOURCE_DIR}")
if(PYNVC_DLPACK_SOURCE_DIR)
    set(dlpack_SOURCE_DIR "${PYNVC_DLPACK_SOURCE_DIR}")
else()
    FetchContent_Populate(
        dlpack
        URL        https://github.com/dmlc/dlpack/archive/refs/tags/v0.8.zip
        URL_HASH   SHA256=80b470d7f5b032b996ce999e470c375cd0a09a20965c168369635b448bb121f4
        SOURCE_DIR dlpack
    )
endif()
message(STATUS "Using dlpack library ${dlpack_SOURCE_DIR}")

# NVTX setup if needed
# Force-enabled for NVTX-instrumented perf investigation (nsys tracing of
# ThreadedDecoder/MultiVideoDecoder). src/CMakeLists.txt normally derives
# USE_NVTX from $ENV{USE_NVTX}; this override makes it unconditional for
# this target regardless of that env var.
set(USE_NVTX 1)
if(NOT DEMUX_ONLY AND USE_NVTX)
    message(STATUS "downloading nvtx library")
    option(FETCHCONTENT_QUIET OFF)
    include(FetchContent)
    fetchcontent_declare(
        nvtx_cpp
        URL https://github.com/NVIDIA/NVTX/archive/v3.1.0.tar.gz
        URL_HASH SHA256=dc4e4a227d04d3da46ad920dfee5f7599ac8d6b2ee1809c9067110fb1cc71ced
        SOURCE_SUBDIR c
        DOWNLOAD_EXTRACT_TIMESTAMP ON
    )
    fetchcontent_makeavailable(nvtx_cpp)
endif()

# Function to create PYD target with specific version
function(create_pyd_target VERSION_SUFFIX VERSION_DEFINE)
    set(TARGET_NAME "${PYTHON_MODULE_PREFIX}${VERSION_SUFFIX}")
    pybind11_add_module(${TARGET_NAME} MODULE 
        ${PY_SOURCES}
    )

    set_property(TARGET ${TARGET_NAME} PROPERTY CXX_STANDARD 17)
    target_include_directories(${TARGET_NAME} 
        PUBLIC 
        ${PY_HDRS}
    )

    target_include_directories(${TARGET_NAME} PRIVATE "${dlpack_SOURCE_DIR}/include")
    
    if(DEMUX_ONLY)
        target_compile_definitions(${TARGET_NAME} PUBLIC -DDEMUX_ONLY=1)
    else()
        if(USE_NVTX)
            target_link_libraries(${TARGET_NAME} PRIVATE nvtx3-cpp)
            target_compile_definitions(${TARGET_NAME} PRIVATE -DUSE_NVTX=1)
        endif()
        target_link_libraries(${TARGET_NAME} PUBLIC VideoCodecSDKUtils${VERSION_SUFFIX})
    endif()

    if(DEMUX_ONLY)
        if(NV_FFMPEG_HAS_BSF)
            target_compile_definitions(${TARGET_NAME} PUBLIC -DHAS_BSF=1)
        endif()
        target_include_directories(${TARGET_NAME} PUBLIC ${NV_FFMPEG_INCLUDE_DIR})
        target_link_libraries(${TARGET_NAME} PUBLIC ${NV_FFMPEG_LIBRARIES})
    endif()

    # Add version-specific define
    target_compile_definitions(${TARGET_NAME} PRIVATE ${VERSION_DEFINE})

    set_target_properties(${TARGET_NAME} PROPERTIES 
        PREFIX ""
        OUTPUT_NAME "${TARGET_NAME}"
        ARCHIVE_OUTPUT_NAME "${TARGET_NAME}"
        PDB_NAME "${TARGET_NAME}"
        INSTALL_RPATH "$ORIGIN"
        BUILD_WITH_INSTALL_RPATH TRUE
    )

    if(DEFINED ENV{SETUPTOOLS_EXT_SUFFIX})
        set_target_properties(${TARGET_NAME} PROPERTIES SUFFIX "$ENV{SETUPTOOLS_EXT_SUFFIX}")
    endif()

    # Emit DT_RPATH (old dtags) instead of DT_RUNPATH so that the bundled
    # ffmpeg next to the module ($ORIGIN) is resolved before any ffmpeg found
    # via LD_LIBRARY_PATH. This mirrors what auditwheel/patchelf --force-rpath
    # does for the released wheels.
    if(UNIX AND NOT APPLE)
        target_link_options(${TARGET_NAME} PRIVATE "LINKER:--disable-new-dtags")
    endif()

    install(TARGETS ${TARGET_NAME}
        LIBRARY DESTINATION ${PYTHON_MODULE_PREFIX}
        RUNTIME DESTINATION ${PYTHON_MODULE_PREFIX}
        ARCHIVE DESTINATION ${PYTHON_MODULE_PREFIX}
    )
endfunction()

# Set architecture-specific flags
if(UNIX)
    # For Linux, ensure 64-bit build on x86_64 (aarch64 is 64-bit by default)
    if(${CMAKE_SYSTEM_PROCESSOR} STREQUAL "x86_64" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "AMD64")
        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -m64")
        set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -m64")
    endif()
endif()

# Create versioned NVENC targets
create_pyd_target("_121" "NVENC_VER_12_1")
create_pyd_target("_130" "NVENC_VER_13_0")

# Create and install VersionCheck module
pybind11_add_module(PyVersionCheck MODULE 
    ${CMAKE_SOURCE_DIR}/src/NvEncVersionCheck/VersionCheck/NvEncVersionCheck.cpp
)

target_include_directories(PyVersionCheck PRIVATE
    ${CMAKE_SOURCE_DIR}/src/NvEncVersionCheck/VersionCheck
)

if(UNIX)
    # Ensure 64-bit build for VersionCheck module on x86_64 (aarch64 is 64-bit by default)
    if(${CMAKE_SYSTEM_PROCESSOR} STREQUAL "x86_64" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "AMD64")
        target_compile_options(PyVersionCheck PRIVATE -m64)
    endif()
endif()

set_target_properties(PyVersionCheck PROPERTIES 
    PREFIX ""
    OUTPUT_NAME "VersionCheck"
    ARCHIVE_OUTPUT_NAME "VersionCheck"
    PDB_NAME "VersionCheck"
    INSTALL_RPATH "$ORIGIN"
    BUILD_WITH_INSTALL_RPATH TRUE
)

if(DEFINED ENV{SETUPTOOLS_EXT_SUFFIX})
    set_target_properties(PyVersionCheck PROPERTIES SUFFIX "$ENV{SETUPTOOLS_EXT_SUFFIX}")
endif()

# For Linux, ensure RPATH is set correctly
if(UNIX)
    set_target_properties(PyVersionCheck PROPERTIES
        INSTALL_RPATH "\$ORIGIN/lib:\$ORIGIN"
        BUILD_WITH_INSTALL_RPATH TRUE
    )
    # Emit DT_RPATH instead of DT_RUNPATH so $ORIGIN takes precedence over
    # LD_LIBRARY_PATH (matches the auditwheel-repaired released wheels).
    if(NOT APPLE)
        target_link_options(PyVersionCheck PRIVATE "LINKER:--disable-new-dtags")
    endif()
endif()

install(TARGETS PyVersionCheck
    LIBRARY DESTINATION ${PYTHON_MODULE_PREFIX}
    RUNTIME DESTINATION ${PYTHON_MODULE_PREFIX}
    ARCHIVE DESTINATION ${PYTHON_MODULE_PREFIX}
)

# Install shared libraries
if(WIN32)
    # Windows DLL installation
    install(
        DIRECTORY "${NV_FFMPEG_RUNTIME_DIR}/"
        DESTINATION PyNvVideoCodec
        FILES_MATCHING 
        PATTERN "*.dll"
    )

    if(NOT DEMUX_ONLY)
        set(_PYNVC_CUDART_DLL_NAMES)
        if(CUDAToolkit_VERSION_MAJOR)
            list(APPEND _PYNVC_CUDART_DLL_NAMES "cudart64_${CUDAToolkit_VERSION_MAJOR}.dll")
        endif()
        list(APPEND _PYNVC_CUDART_DLL_NAMES cudart64_13.dll cudart64_12.dll cudart64_11.dll)
        list(REMOVE_DUPLICATES _PYNVC_CUDART_DLL_NAMES)

        set(_PYNVC_CUDART_HINTS)
        if(CMAKE_SYSTEM_PROCESSOR STREQUAL "ARM64")
            if(CUDAToolkit_ROOT)
                list(APPEND _PYNVC_CUDART_HINTS "${CUDAToolkit_ROOT}/bin/arm64")
            endif()
            if(DEFINED ENV{CUDA_PATH} AND NOT "$ENV{CUDA_PATH}" STREQUAL "")
                list(APPEND _PYNVC_CUDART_HINTS "$ENV{CUDA_PATH}/bin/arm64")
            endif()
        else()
            if(CUDAToolkit_BIN_DIR)
                list(APPEND _PYNVC_CUDART_HINTS "${CUDAToolkit_BIN_DIR}")
            endif()
            if(CUDAToolkit_ROOT)
                list(APPEND _PYNVC_CUDART_HINTS "${CUDAToolkit_ROOT}/bin")
            endif()
            if(DEFINED ENV{CUDA_PATH} AND NOT "$ENV{CUDA_PATH}" STREQUAL "")
                list(APPEND _PYNVC_CUDART_HINTS "$ENV{CUDA_PATH}/bin")
            endif()
        endif()
        list(REMOVE_DUPLICATES _PYNVC_CUDART_HINTS)

        find_file(
            PYNVC_CUDA_RUNTIME_DLL
            NAMES ${_PYNVC_CUDART_DLL_NAMES}
            HINTS ${_PYNVC_CUDART_HINTS}
            NO_DEFAULT_PATH
        )
        if(NOT PYNVC_CUDA_RUNTIME_DLL AND NOT CMAKE_SYSTEM_PROCESSOR STREQUAL "ARM64")
            find_file(
                PYNVC_CUDA_RUNTIME_DLL
                NAMES ${_PYNVC_CUDART_DLL_NAMES}
            )
        endif()

        if(NOT PYNVC_CUDA_RUNTIME_DLL)
            message(FATAL_ERROR "CUDA runtime DLL was not found; cannot build a self-contained Windows wheel")
        endif()

        install(FILES "${PYNVC_CUDA_RUNTIME_DLL}" DESTINATION ${PYTHON_MODULE_PREFIX})
    endif()
else()
    # Linux .so installation
    install(
        DIRECTORY "${NV_FFMPEG_LIBRARY_DIR}/"
        DESTINATION ${PYTHON_MODULE_PREFIX}
        FILES_MATCHING 
        PATTERN "*.so*"
        PATTERN "*.a" EXCLUDE
    )
endif()
