# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# SPDX-License-Identifier: Apache-2.0

# ################################################################################
#
# This example demonstrates StridedMemoryView for interfacing with foreign
# C/C++ functions, using JIT-compiled CPU code via cffi. Requires cffi.
#
# ################################################################################

# /// script
# dependencies = ["cuda_bindings", "cuda_core", "cffi", "setuptools"]
# ///

import importlib
import string
import sys
import tempfile
from contextlib import contextmanager

try:
    from cffi import FFI
except ImportError:
    print("cffi is not installed, this example requires cffi", file=sys.stderr)
    sys.exit(1)
import numpy as np

from cuda.core.utils import StridedMemoryView, args_viewable_as_strided_memory

# ################################################################################
#
# Usually this entire code block is in a separate file, built as a Python extension
# module that can be imported by users at run time. For illustrative purposes we
# use JIT compilation to make this demo self-contained.
#
# Here we assume an in-place operation, equivalent to the following NumPy code:
#
#   >>> arr = ...
#   >>> assert arr.dtype == np.int32
#   >>> assert arr.ndim == 1
#   >>> arr += np.arange(arr.size, dtype=arr.dtype)
#
# is implemented for both CPU and GPU at low-level, with the following C function
# signature:
func_name = "inplace_plus_arange_N"
func_sig = f"void {func_name}(int* data, size_t N)"


# Now we are prepared to run the code from the user's perspective!
#
# ################################################################################


# Below, as a user we want to perform the said in-place operation on a CPU
# or GPU, by calling the corresponding function implemented "elsewhere"
# (in the body of run function).


# We assume the 0-th argument supports either DLPack or CUDA Array Interface (both
# of which are supported by StridedMemoryView).
@args_viewable_as_strided_memory((0,))
def my_func(arr, cpu_prog, cpu_func):
    # Create a memory view over arr (assumed to be a 1D array of int32). The stream
    # ordering is taken care of, so that arr can be safely accessed on our work
    # stream (ordered after a data stream on which arr is potentially prepared).
    view = arr.view(-1)
    assert isinstance(view, StridedMemoryView)
    assert len(view.shape) == 1
    assert view.dtype == np.int32
    assert not view.is_device_accessible

    size = view.shape[0]
    # DLPack also supports host arrays. We want to know if the array data is
    # accessible from the GPU, and dispatch to the right routine accordingly.
    cpu_func(cpu_prog.cast("int*", view.ptr), size)


def _create_cpu_program():
    # Here is a concrete (very naive!) implementation on CPU:
    cpu_code = string.Template(r"""
    extern "C"
    $func_sig {
        for (size_t i = 0; i < N; i++) {
            data[i] += i;
        }
    }
    """).substitute(func_sig=func_sig)
    # This is cffi's way of JIT compiling & loading a CPU function. cffi builds an
    # extension module that has the Python binding to the underlying C function.
    # For more details, please refer to cffi's documentation.
    cpu_prog = FFI()
    cpu_prog.cdef(f"{func_sig};")
    cpu_prog.set_source(
        "_cpu_obj",
        cpu_code,
        source_extension=".cpp",
        extra_compile_args=["-std=c++11"],
    )
    return cpu_prog


@contextmanager
def _compiled_cpu_func(cpu_prog, temp_dir):
    saved_sys_path = sys.path.copy()
    try:
        cpu_prog.compile(tmpdir=temp_dir)
        sys.path.append(temp_dir)
        cpu_func = getattr(importlib.import_module("_cpu_obj.lib"), func_name)
        yield cpu_func
    finally:
        sys.path = saved_sys_path
        # Ensure cffi modules are unloadable before removing the temp build dir.
        sys.modules.pop("_cpu_obj.lib", None)
        sys.modules.pop("_cpu_obj", None)


def _run_example(cpu_prog, cpu_func):
    arr_cpu = np.zeros(1024, dtype=np.int32)
    print(f"before: {arr_cpu[:10]=}", file=sys.stderr)
    my_func(arr_cpu, cpu_prog, cpu_func)
    print(f"after: {arr_cpu[:10]=}", file=sys.stderr)
    assert np.allclose(arr_cpu, np.arange(1024, dtype=np.int32))


def main():
    cpu_prog = _create_cpu_program()
    with tempfile.TemporaryDirectory() as temp_dir, _compiled_cpu_func(cpu_prog, temp_dir) as cpu_func:
        _run_example(cpu_prog, cpu_func)


if __name__ == "__main__":
    main()
