// Copyright (C) 2018-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

/**
 * @brief a header that defines wrappers for internal NPU plugin-specific
 * LevelZero context and LevelZero shared memory tensors
 *
 * @file openvino/runtime/intel_npu/level_zero/level_zero.hpp
 */
#pragma once

#include "openvino/runtime/core.hpp"
#include "openvino/runtime/intel_npu/properties.hpp"
#include "openvino/runtime/intel_npu/remote_properties.hpp"
#include "openvino/runtime/remote_context.hpp"
#include "openvino/runtime/remote_tensor.hpp"

namespace ov {
namespace intel_npu {

/**
 * @defgroup ov_runtime_level_zero_npu_cpp_api Intel NPU LevelZero interoperability
 * @ingroup ov_runtime_cpp_api
 * Set of C++ classes and properties to work with Remote API for Intel NPU LevelZero plugin.
 */

/**
 * @brief Namespace with Intel NPU LevelZero specific remote objects
 */
namespace level_zero {

/**
 * @brief This class represents an abstraction for NPU plugin remote tensor
 * which can be shared with user-supplied LevelZero buffer.
 * The plugin object derived from this class can be obtained with ZeroContext::create_tensor() call.
 * @note User can obtain Level Zero buffer handle from this class.
 * @ingroup ov_runtime_level_zero_npu_cpp_api
 */
class ZeroBufferTensor : public RemoteTensor {
public:
    /**
     * @brief Checks that type defined runtime parameters are presented in remote object
     * @param tensor a tensor to check
     */
    static void type_check(const Tensor& tensor) {
        RemoteTensor::type_check(tensor,
                                 {{std::string(ov::intel_npu::mem_handle.name()), {}},
                                  {std::string(ov::intel_npu::mem_type.name()),
                                   {ov::Any(ov::intel_npu::MemType::L0_INTERNAL_BUF).as<std::string>(),
                                    ov::Any(ov::intel_npu::MemType::SHARED_BUF).as<std::string>(),
                                    ov::Any(ov::intel_npu::MemType::MMAPED_FILE).as<std::string>(),
                                    ov::Any(ov::intel_npu::MemType::CPU_VA).as<std::string>()}}});
    }

    /**
     * @brief Returns the underlying LevelZero memory object handle.
     * @return underlying void* memory object handle
     */
    void* get() {
        return get_params().at(ov::intel_npu::mem_handle.name()).as<void*>();
    }
};

/**
 * @brief This class represents an abstraction for NPU plugin remote context
 * which is shared with LevelZero context object.
 * The plugin object derived from this class can be obtained either with
 * CompiledModel::get_context() or Core::create_context() calls.
 * @ingroup ov_runtime_level_zero_npu_cpp_api
 */
class ZeroContext : public RemoteContext {
protected:
    /**
     * @brief NPU device name
     */
    static constexpr const char* device_name = "NPU";

    /**
     * @brief Default constructor which can be used in derived classes to avoid multiple create_context() calls
     */
    ZeroContext() = default;

public:
    // Needed to make create_tensor overloads from base class visible for user
    using RemoteContext::create_tensor;

    /**
     * @brief Constructs context object from user-supplied LevelZero context handle
     * @param core A reference to OpenVINO Runtime Core object
     */
    ZeroContext(Core& core) {
        *this = core.get_default_context(device_name).as<ZeroContext>();
    }

    /**
     * @brief Returns the underlying LevelZero context handle.
     * @return `void*`
     */
    void* get() {
        return get_params().at(l0_context.name()).as<void*>();
    }

    /**
     * @brief This function is used to obtain remote tensor object from user-supplied NT handle object or raw pointer
     * @param type Tensor element type
     * @param shape Tensor shape
     * @param buffer A void* object that should be wrapped by a remote tensor
     * @param memory_type Memory type to use (default: SHARED_BUF)
     * @param tensor_type Type of the tensor to be shared, input, output or binded (default: BINDED)
     * @return A remote tensor instance
     */
    ZeroBufferTensor create_tensor(const element::Type type,
                                   const Shape& shape,
                                   void* buffer,
                                   const MemType memory_type = MemType::SHARED_BUF,
                                   const TensorType tensor_type = TensorType::BINDED) {
        OPENVINO_ASSERT(memory_type == MemType::SHARED_BUF || memory_type == MemType::CPU_VA,
                        "Only SHARED_BUF and CPU_VA memory types are supported for raw buffer pointer or NT handle");

        AnyMap params = {{mem_type.name(), memory_type},
                         {mem_handle.name(), buffer},
                         {ov::intel_npu::tensor_type.name(), tensor_type}};
        return create_tensor(type, shape, params).as<ZeroBufferTensor>();
    }

    /**
     * @brief This function is used to obtain remote tensor object from user-supplied DMA-BUF System Heap object
     * @param type Tensor element type
     * @param shape Tensor shape
     * @param fd A int object that should be wrapped by a remote tensor
     * @param memory_type Memory type to use (default: SHARED_BUF)
     * @param tensor_type Type of the tensor to be shared, input, output or binded (default: BINDED)
     * @return A remote tensor instance
     */
    ZeroBufferTensor create_tensor(const element::Type type,
                                   const Shape& shape,
                                   int fd,
                                   const MemType memory_type = MemType::SHARED_BUF,
                                   const TensorType tensor_type = TensorType::BINDED) {
        OPENVINO_ASSERT(memory_type == MemType::SHARED_BUF,
                        "Only SHARED_BUF memory type is supported for DMA-BUF file descriptor");

        AnyMap params = {{mem_type.name(), memory_type},
                         {mem_handle.name(), reinterpret_cast<void*>(static_cast<intptr_t>(fd))},
                         {ov::intel_npu::tensor_type.name(), tensor_type}};
        return create_tensor(type, shape, params).as<ZeroBufferTensor>();
    }

    /**
     * @brief This function is used to obtain remote tensor object from a file
     * @param type Tensor element type
     * @param shape Tensor shape
     * @param file_name A string object that should be wrapped by a remote tensor
     * @param offset_in_bytes Offset in bytes from the beginning of the file
     * @param tensor_type Type of the tensor to be shared, input, output or binded
     * @return A remote tensor instance
     */
    ZeroBufferTensor create_tensor(const element::Type type,
                                   const Shape& shape,
                                   const FileDescriptor& file_descriptor,
                                   const TensorType tensor_type = TensorType::INPUT) {
        AnyMap params = {{ov::intel_npu::mem_type.name(), MemType::MMAPED_FILE},
                         {ov::intel_npu::file_descriptor.name(), file_descriptor},
                         {ov::intel_npu::tensor_type.name(), tensor_type}};
        return create_tensor(type, shape, params).as<ZeroBufferTensor>();
    }

    /**
     * @brief This function is used to obtain remote tensor object
     * @param type Tensor element type
     * @param shape Tensor shape
     * @param tensor_type Type of the tensor to be shared, input, output or binded
     * @return A remote tensor instance
     */
    ZeroBufferTensor create_l0_host_tensor(const element::Type type,
                                           const Shape& shape,
                                           const TensorType tensor_type = TensorType::BINDED) {
        AnyMap params = {{mem_type.name(), MemType::L0_INTERNAL_BUF}, {ov::intel_npu::tensor_type.name(), tensor_type}};
        return create_tensor(type, shape, params).as<ZeroBufferTensor>();
    }
};

}  // namespace level_zero
}  // namespace intel_npu
}  // namespace ov
