/*
 * Copyright (C) 2022-2026 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

#include "level_zero_driver/api/ext/ze_graph.hpp"

#include "level_zero_driver/api/trace/trace_ze_graph_ext.hpp"
#include "level_zero_driver/api/trace/trace_ze_graph_profiling_ext.hpp"
#include "level_zero_driver/api/zet_misc.hpp"
#include "level_zero_driver/include/l0_exception.hpp"
#include "level_zero_driver/source/cmdlist.hpp"
#include "level_zero_driver/source/context.hpp"
#include "level_zero_driver/source/device.hpp"
#include "level_zero_driver/source/driver.hpp"
#include "level_zero_driver/source/event.hpp"
#include "level_zero_driver/source/ext/graph.hpp"
#include "level_zero_driver/source/ext/profiling_data.hpp"
#include "level_zero_driver/source/ext/query_network.hpp"
#include "vpu_driver/source/utilities/log.hpp"

#include <algorithm>
#include <loader/ze_loader.h>
#include <memory>
#include <vector>
#include <ze_api.h>
#include <ze_graph_ext.h>
#include <ze_graph_profiling_ext.h>

namespace L0 {

template <typename T>
const char *getGraphExtStructureName() {
    if constexpr (std::is_same_v<::ze_device_graph_properties_t, T>)
        return "ze_device_graph_properties_t";
    if constexpr (std::is_same_v<::ze_device_graph_properties_2_t, T>)
        return "ze_device_graph_properties_2_t";
    if constexpr (std::is_same_v<::ze_graph_desc_t, T>)
        return "ze_graph_desc_t";
    if constexpr (std::is_same_v<::ze_graph_desc_2_t, T>)
        return "ze_graph_desc_2_t";
    if constexpr (std::is_same_v<::ze_graph_properties_t, T>)
        return "ze_graph_properties_t";
    if constexpr (std::is_same_v<::ze_graph_properties_2_t, T>)
        return "ze_graph_properties_2_t";
    if constexpr (std::is_same_v<::ze_graph_properties_3_t, T>)
        return "ze_graph_properties_3_t";
    if constexpr (std::is_same_v<::ze_graph_argument_properties_t, T>)
        return "ze_graph_argument_properties_t";
    if constexpr (std::is_same_v<::ze_graph_argument_properties_2_t, T>)
        return "ze_graph_argument_properties_2_t";
    if constexpr (std::is_same_v<::ze_graph_argument_properties_3_t, T>)
        return "ze_graph_argument_properties_3_t";
    if constexpr (std::is_same_v<::ze_graph_argument_metadata_t, T>)
        return "ze_graph_argument_metadata_t";
    if constexpr (std::is_same_v<::ze_graph_input_hash_t, T>)
        return "ze_graph_input_hash_t";
    if constexpr (std::is_same_v<::ze_graph_argument_property_strides_t, T>)
        return "ze_graph_argument_property_strides_t";
    if constexpr (std::is_same_v<::ze_graph_argument_value_tensor_t, T>)
        return "ze_graph_argument_value_tensor_t";
    if constexpr (std::is_same_v<::ze_graph_argument_value_strides_t, T>)
        return "ze_graph_argument_value_strides_t";
    return "unknown_graph_structure";
}

template <typename T>
ze_structure_type_graph_ext_t getGraphExtStructureType() {
    if constexpr (std::is_same_v<::ze_device_graph_properties_t, T>)
        return ZE_STRUCTURE_TYPE_DEVICE_GRAPH_PROPERTIES;
    if constexpr (std::is_same_v<::ze_device_graph_properties_2_t, T>)
        return ZE_STRUCTURE_TYPE_DEVICE_GRAPH_PROPERTIES_2;
    if constexpr (std::is_same_v<::ze_graph_desc_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_DESC;
    if constexpr (std::is_same_v<::ze_graph_desc_2_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_DESC_2;
    if constexpr (std::is_same_v<::ze_graph_properties_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_PROPERTIES;
    if constexpr (std::is_same_v<::ze_graph_properties_2_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_PROPERTIES_2;
    if constexpr (std::is_same_v<::ze_graph_properties_3_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_PROPERTIES_3;
    if constexpr (std::is_same_v<::ze_graph_argument_properties_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_PROPERTIES;
    if constexpr (std::is_same_v<::ze_graph_argument_properties_2_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_PROPERTIES_2;
    if constexpr (std::is_same_v<::ze_graph_argument_properties_3_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_PROPERTIES_3;
    if constexpr (std::is_same_v<::ze_graph_argument_metadata_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_METADATA;
    if constexpr (std::is_same_v<::ze_graph_input_hash_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_INPUT_HASH;
    if constexpr (std::is_same_v<::ze_graph_argument_property_strides_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_PROPERTY_STRIDES;
    if constexpr (std::is_same_v<::ze_graph_argument_value_tensor_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_TENSOR;
    if constexpr (std::is_same_v<::ze_graph_argument_value_strides_t, T>)
        return ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_STRIDES;
    return ZE_STRUCTURE_TYPE_GRAPH_FORCE_UINT32;
}

template <typename T>
ze_result_t validateGraphStructure(const T *pStruct,
                                   std::vector<ze_structure_type_graph_ext_t> validPNext = {}) {
    if (pStruct == nullptr) {
        LOG_E("Invalid %s structure pointer", getGraphExtStructureName<T>());
        return ZE_RESULT_ERROR_INVALID_NULL_POINTER;
    }
    if (!Driver::getInstance()->getEnvVariables().extensionValidation) {
        return ZE_RESULT_SUCCESS;
    }
    if (pStruct->stype != getGraphExtStructureType<T>()) {
        LOG_E("Invalid %s structure stype", getGraphExtStructureName<T>());
        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
    }
    for (const auto *pNext = pStruct->pNext; pNext != nullptr;) {
        auto *stype = reinterpret_cast<const ze_structure_type_graph_ext_t *>(pNext);
        if (!std::any_of(
                validPNext.begin(),
                validPNext.end(),
                [stype](ze_structure_type_graph_ext_t validType) { return *stype == validType; })) {
            LOG_E("Unsupported %x structure in pNext chain of %s",
                  *stype,
                  getGraphExtStructureName<T>());
            return ZE_RESULT_ERROR_INVALID_ARGUMENT;
        }

        pNext = reinterpret_cast<const ze_base_desc_t *>(pNext)->pNext;
    }
    return ZE_RESULT_SUCCESS;
}

ze_result_t ZE_APICALL zeGraphCreate(ze_context_handle_t hContext,
                                     ze_device_handle_t hDevice,
                                     const ze_graph_desc_t *pDesc,
                                     ze_graph_handle_t *phGraph) {
    trace_zeGraphCreate(hContext, hDevice, pDesc, phGraph);
    ze_graph_desc_2_t desc{};
    ze_result_t ret;

    ret = validateGraphStructure(pDesc, {ZE_STRUCTURE_TYPE_GRAPH_INPUT_HASH});
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_CONTEXT, hContext);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    desc.stype = pDesc->stype;
    desc.pNext = pDesc->pNext;
    desc.format = pDesc->format;
    desc.inputSize = pDesc->inputSize;
    desc.pInput = pDesc->pInput;
    desc.pBuildFlags = pDesc->pBuildFlags;
    desc.flags = 0;

    L0_HANDLE_EXCEPTION(ret, L0::Graph::create(hContext, hDevice, &desc, phGraph));

exit:
    trace_zeGraphCreate(ret, hContext, hDevice, pDesc, phGraph);
    return ret;
}

ze_result_t ZE_APICALL zeGraphDestroy(ze_graph_handle_t hGraph) {
    trace_zeGraphDestroy(hGraph);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::fromHandle(hGraph)->destroy());

exit:
    trace_zeGraphDestroy(ret, hGraph);
    return ret;
}

ze_result_t ZE_APICALL zeGraphGetProperties(ze_graph_handle_t hGraph,
                                            ze_graph_properties_t *pGraphProperties) {
    trace_zeGraphGetProperties(hGraph, pGraphProperties);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(pGraphProperties);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::fromHandle(hGraph)->getProperties(pGraphProperties));

exit:
    trace_zeGraphGetProperties(ret, hGraph, pGraphProperties);
    return ret;
}

ze_result_t ZE_APICALL zeGraphGetProperties2(ze_graph_handle_t hGraph,
                                             ze_graph_properties_2_t *pGraphProperties) {
    trace_zeGraphGetProperties2(hGraph, pGraphProperties);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(pGraphProperties);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::fromHandle(hGraph)->getProperties2(pGraphProperties));

exit:
    trace_zeGraphGetProperties2(ret, hGraph, pGraphProperties);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphGetArgumentProperties(ze_graph_handle_t hGraph,
                             uint32_t argIndex,
                             ze_graph_argument_properties_t *pGraphArgumentProperties) {
    trace_zeGraphGetArgumentProperties(hGraph, argIndex, pGraphArgumentProperties);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(pGraphArgumentProperties,
                                 {ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_PROPERTY_STRIDES});
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(
        ret,
        L0::Graph::fromHandle(hGraph)->getArgumentProperties(argIndex, pGraphArgumentProperties));

exit:
    trace_zeGraphGetArgumentProperties(ret, hGraph, argIndex, pGraphArgumentProperties);
    return ret;
}

ze_result_t ZE_APICALL zeGraphSetArgumentValue(ze_graph_handle_t hGraph,
                                               uint32_t argIndex,
                                               const void *pArgValue) {
    trace_zeGraphSetArgumentValue(hGraph, argIndex, pArgValue);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::fromHandle(hGraph)->setArgumentValue(argIndex, pArgValue));

exit:
    trace_zeGraphSetArgumentValue(ret, hGraph, argIndex, pArgValue);
    return ret;
}

ze_result_t ZE_APICALL zeGraphSetArgumentValue2(ze_graph_handle_t hGraph,
                                                uint32_t argIndex,
                                                const void *pArgValue) {
    trace_zeGraphSetArgumentValue2(hGraph, argIndex, pArgValue);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::fromHandle(hGraph)->setArgumentValue2(argIndex, pArgValue));

exit:
    trace_zeGraphSetArgumentValue2(ret, hGraph, argIndex, pArgValue);
    return ret;
}

static ze_result_t appendGraphInitialize(ze_command_list_handle_t hCommandList,
                                         ze_graph_handle_t hGraph,
                                         ze_event_handle_t hSignalEvent,
                                         uint32_t numWaitEvents,
                                         ze_event_handle_t *phWaitEvents) {
    auto result = translateHandle(ZEL_HANDLE_COMMAND_LIST, hCommandList);
    if (result != ZE_RESULT_SUCCESS) {
        return result;
    }

    if (hSignalEvent != nullptr) {
        auto result = translateHandle(ZEL_HANDLE_EVENT, hSignalEvent);
        if (result != ZE_RESULT_SUCCESS)
            return result;
    }

    std::unique_ptr<ze_event_handle_t[]> phWaitEventsTranslated = nullptr;
    if (numWaitEvents > 0) {
        if (phWaitEvents == nullptr)
            return ZE_RESULT_ERROR_INVALID_SIZE;

        phWaitEventsTranslated =
            std::unique_ptr<ze_event_handle_t[]>(new ze_event_handle_t[numWaitEvents]);
        for (uint32_t i = 0; i < numWaitEvents; i++) {
            auto result =
                translateHandle(ZEL_HANDLE_EVENT, phWaitEvents[i], &phWaitEventsTranslated[i]);
            if (result != ZE_RESULT_SUCCESS)
                return result;
        }
    }

    return L0::CommandList::fromHandle(hCommandList)
        ->appendGraphInitialize(hGraph, hSignalEvent, numWaitEvents, phWaitEventsTranslated.get());
}

ze_result_t ZE_APICALL zeAppendGraphInitialize(ze_command_list_handle_t hCommandList,
                                               ze_graph_handle_t hGraph,
                                               ze_event_handle_t hSignalEvent,
                                               uint32_t numWaitEvents,
                                               ze_event_handle_t *phWaitEvents) {
    trace_zeAppendGraphInitialize(hCommandList, hGraph, hSignalEvent, numWaitEvents, phWaitEvents);
    ze_result_t ret;

    if (hCommandList == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(
        ret,
        appendGraphInitialize(hCommandList, hGraph, hSignalEvent, numWaitEvents, phWaitEvents));

exit:
    trace_zeAppendGraphInitialize(ret,
                                  hCommandList,
                                  hGraph,
                                  hSignalEvent,
                                  numWaitEvents,
                                  phWaitEvents);
    return ret;
}

ze_result_t ZE_APICALL zeGraphInitialize(ze_graph_handle_t hGraph) {
    trace_zeGraphInitialize(hGraph);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::fromHandle(hGraph)->parserInitialize());

exit:
    trace_zeGraphInitialize(ret, hGraph);
    return ret;
}

static ze_result_t appendGraphExecute(ze_command_list_handle_t hCommandList,
                                      ze_graph_handle_t hGraph,
                                      ze_graph_profiling_query_handle_t hProfilingQuery,
                                      ze_event_handle_t hSignalEvent,
                                      uint32_t numWaitEvents,
                                      ze_event_handle_t *phWaitEvents) {
    auto result = translateHandle(ZEL_HANDLE_COMMAND_LIST, hCommandList);
    if (result != ZE_RESULT_SUCCESS) {
        return result;
    }

    if (hSignalEvent != nullptr) {
        auto result = translateHandle(ZEL_HANDLE_EVENT, hSignalEvent);
        if (result != ZE_RESULT_SUCCESS)
            return result;
    }

    std::unique_ptr<ze_event_handle_t[]> phWaitEventsTranslated = nullptr;
    if (numWaitEvents > 0) {
        if (phWaitEvents == nullptr)
            return ZE_RESULT_ERROR_INVALID_SIZE;

        phWaitEventsTranslated =
            std::unique_ptr<ze_event_handle_t[]>(new ze_event_handle_t[numWaitEvents]);
        for (uint32_t i = 0; i < numWaitEvents; i++) {
            auto result =
                translateHandle(ZEL_HANDLE_EVENT, phWaitEvents[i], &phWaitEventsTranslated[i]);
            if (result != ZE_RESULT_SUCCESS)
                return result;
        }
    }

    return L0::CommandList::fromHandle(hCommandList)
        ->appendGraphExecute(hGraph,
                             hProfilingQuery,
                             hSignalEvent,
                             numWaitEvents,
                             phWaitEventsTranslated.get());
}

ze_result_t ZE_APICALL zeAppendGraphExecute(ze_command_list_handle_t hCommandList,
                                            ze_graph_handle_t hGraph,
                                            ze_graph_profiling_query_handle_t hProfilingQuery,
                                            ze_event_handle_t hSignalEvent,
                                            uint32_t numWaitEvents,
                                            ze_event_handle_t *phWaitEvents) {
    trace_zeAppendGraphExecute(hCommandList,
                               hGraph,
                               hProfilingQuery,
                               hSignalEvent,
                               numWaitEvents,
                               phWaitEvents);
    ze_result_t ret;

    if (hCommandList == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret,
                        appendGraphExecute(hCommandList,
                                           hGraph,
                                           hProfilingQuery,
                                           hSignalEvent,
                                           numWaitEvents,
                                           phWaitEvents));

exit:
    trace_zeAppendGraphExecute(ret,
                               hCommandList,
                               hGraph,
                               hProfilingQuery,
                               hSignalEvent,
                               numWaitEvents,
                               phWaitEvents);
    return ret;
}

ze_result_t ZE_APICALL zeGraphGetNativeBinary(ze_graph_handle_t hGraph,
                                              size_t *pSize,
                                              uint8_t *pGraphNativeBinary) {
    trace_zeGraphGetNativeBinary(hGraph, pSize, pGraphNativeBinary);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret,
                        L0::Graph::fromHandle(hGraph)->getNativeBinary(pSize, pGraphNativeBinary));

exit:
    trace_zeGraphGetNativeBinary(ret, hGraph, pSize, pGraphNativeBinary);
    return ret;
}

ze_result_t ZE_APICALL zeGraphGetNativeBinary2(ze_graph_handle_t hGraph,
                                               size_t *pSize,
                                               const uint8_t **pGraphNativeBinary) {
    trace_zeGraphGetNativeBinary2(hGraph, pSize, pGraphNativeBinary);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    if (pSize == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_POINTER;
        goto exit;
    }

    if (pGraphNativeBinary == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_POINTER;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret,
                        L0::Graph::fromHandle(hGraph)->getNativeBinary2(pSize, pGraphNativeBinary));

exit:
    trace_zeGraphGetNativeBinary2(ret, hGraph, pSize, pGraphNativeBinary);
    return ret;
}

ze_result_t ZE_APICALL
zeDeviceGetGraphProperties(ze_device_handle_t hDevice,
                           ze_device_graph_properties_t *pDeviceGraphProperties) {
    trace_zeDeviceGetGraphProperties(hDevice, pDeviceGraphProperties);
    ze_result_t ret;

    if (hDevice == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = validateGraphStructure(pDeviceGraphProperties);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::getDeviceGraphProperties(hDevice, pDeviceGraphProperties));

exit:
    trace_zeDeviceGetGraphProperties(ret, hDevice, pDeviceGraphProperties);
    return ret;
}

ze_result_t ZE_APICALL
zeDeviceGetGraphProperties2(ze_device_handle_t hDevice,
                            ze_device_graph_properties_2_t *pDeviceGraphProperties2) {
    trace_zeDeviceGetGraphProperties2(hDevice, pDeviceGraphProperties2);
    ze_result_t ret;

    if (hDevice == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(pDeviceGraphProperties2);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret,
                        L0::Graph::getDeviceGraphProperties2(hDevice, pDeviceGraphProperties2));

exit:
    trace_zeDeviceGetGraphProperties2(ret, hDevice, pDeviceGraphProperties2);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphGetArgumentMetadata(ze_graph_handle_t hGraph,
                           uint32_t argIndex,
                           ze_graph_argument_metadata_t *pGraphArgumentMetadata) {
    trace_zeGraphGetArgumentMetadata(hGraph, argIndex, pGraphArgumentMetadata);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(pGraphArgumentMetadata);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(
        ret,
        L0::Graph::fromHandle(hGraph)->getArgumentMetadata(argIndex, pGraphArgumentMetadata));

exit:
    trace_zeGraphGetArgumentMetadata(ret, hGraph, argIndex, pGraphArgumentMetadata);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphGetArgumentProperties2(ze_graph_handle_t hGraph,
                              uint32_t argIndex,
                              ze_graph_argument_properties_2_t *pGraphArgumentProperties) {
    trace_zeGraphGetArgumentProperties2(hGraph, argIndex, pGraphArgumentProperties);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(pGraphArgumentProperties,
                                 {ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_PROPERTY_STRIDES});
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(
        ret,
        L0::Graph::fromHandle(hGraph)->getArgumentProperties2(argIndex, pGraphArgumentProperties));

exit:
    trace_zeGraphGetArgumentProperties2(ret, hGraph, argIndex, pGraphArgumentProperties);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphGetArgumentProperties3(ze_graph_handle_t hGraph,
                              uint32_t argIndex,
                              ze_graph_argument_properties_3_t *pGraphArgumentProperties) {
    trace_zeGraphGetArgumentProperties3(hGraph, argIndex, pGraphArgumentProperties);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(pGraphArgumentProperties,
                                 {ZE_STRUCTURE_TYPE_GRAPH_ARGUMENT_PROPERTY_STRIDES});
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(
        ret,
        L0::Graph::fromHandle(hGraph)->getArgumentProperties3(argIndex, pGraphArgumentProperties));

exit:
    trace_zeGraphGetArgumentProperties3(ret, hGraph, argIndex, pGraphArgumentProperties);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphQueryNetworkCreate(ze_context_handle_t hContext,
                          ze_device_handle_t hDevice,
                          const ze_graph_desc_t *pDesc,
                          ze_graph_query_network_handle_t *phGraphQueryNetwork) {
    trace_zeGraphQueryNetworkCreate(hContext, hDevice, pDesc, phGraphQueryNetwork);
    ze_graph_desc_2_t desc{};
    ze_result_t ret;

    ret = validateGraphStructure(pDesc, {ZE_STRUCTURE_TYPE_GRAPH_INPUT_HASH});
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    desc.stype = pDesc->stype;
    desc.pNext = pDesc->pNext;
    desc.format = pDesc->format;
    desc.inputSize = pDesc->inputSize;
    desc.pInput = pDesc->pInput;
    desc.pBuildFlags = pDesc->pBuildFlags;
    desc.flags = 0;

    ret = zeGraphQueryNetworkCreate2(hContext, hDevice, &desc, phGraphQueryNetwork);

exit:
    trace_zeGraphQueryNetworkCreate(ret, hContext, hDevice, pDesc, phGraphQueryNetwork);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphQueryNetworkDestroy(ze_graph_query_network_handle_t hGraphQueryNetwork) {
    trace_zeGraphQueryNetworkDestroy(hGraphQueryNetwork);
    ze_result_t ret;

    if (hGraphQueryNetwork == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::QueryNetwork::fromHandle(hGraphQueryNetwork)->destroy());

exit:
    trace_zeGraphQueryNetworkDestroy(ret, hGraphQueryNetwork);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphQueryNetworkGetSupportedLayers(ze_graph_query_network_handle_t hGraphQueryNetwork,
                                      size_t *pSize,
                                      char *pSupportedLayers) {
    trace_zeGraphQueryNetworkGetSupportedLayers(hGraphQueryNetwork, pSize, pSupportedLayers);
    ze_result_t ret;

    if (hGraphQueryNetwork == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret,
                        L0::QueryNetwork::fromHandle(hGraphQueryNetwork)
                            ->getSupportedLayers(pSize, pSupportedLayers));

exit:
    trace_zeGraphQueryNetworkGetSupportedLayers(ret, hGraphQueryNetwork, pSize, pSupportedLayers);
    return ret;
}

ze_result_t ZE_APICALL zeGraphBuildLogGetString(ze_graph_handle_t hGraph,
                                                uint32_t *pSize,
                                                char *pBuildLog) {
    trace_zeGraphBuildLogGetString(hGraph, pSize, pBuildLog);
    ze_result_t ret;

    if (hGraph != nullptr) {
        // Case when the graph creation was successful, no logs from compiler
        ret = ZE_RESULT_SUCCESS;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::getLogString(pSize, pBuildLog));

exit:
    trace_zeGraphBuildLogGetString(ret, hGraph, pSize, pBuildLog);
    return ret;
}

ze_result_t ZE_APICALL zeGraphCreate2(ze_context_handle_t hContext,
                                      ze_device_handle_t hDevice,
                                      const ze_graph_desc_2_t *desc,
                                      ze_graph_handle_t *phGraph) {
    trace_zeGraphCreate2(hContext, hDevice, desc, phGraph);
    ze_result_t ret;

    if (hDevice == nullptr || hContext == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(desc, {ZE_STRUCTURE_TYPE_GRAPH_INPUT_HASH});
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_CONTEXT, hContext);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::create(hContext, hDevice, desc, phGraph));

exit:
    trace_zeGraphCreate2(ret, hContext, hDevice, desc, phGraph);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphQueryNetworkCreate2(ze_context_handle_t hContext,
                           ze_device_handle_t hDevice,
                           const ze_graph_desc_2_t *desc,
                           ze_graph_query_network_handle_t *phGraphQueryNetwork) {
    trace_zeGraphQueryNetworkCreate2(hContext, hDevice, desc, phGraphQueryNetwork);
    ze_result_t ret;

    if (hDevice == nullptr || hContext == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(desc, {ZE_STRUCTURE_TYPE_GRAPH_INPUT_HASH});
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_CONTEXT, hContext);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret,
                        L0::QueryNetwork::create(hContext, hDevice, desc, phGraphQueryNetwork));

exit:
    trace_zeGraphQueryNetworkCreate2(ret, hContext, hDevice, desc, phGraphQueryNetwork);
    return ret;
}

ze_result_t ZE_APICALL zeGraphQueryContextMemory(ze_context_handle_t hContext,
                                                 ze_graph_memory_query_type_t type,
                                                 ze_graph_memory_query_t *query) {
    trace_zeGraphQueryContextMemory(hContext, type, query);
    ze_result_t ret;

    if (hContext == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_CONTEXT, hContext);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Context::fromHandle(hContext)->queryContextMemory(type, query));

exit:
    trace_zeGraphQueryContextMemory(ret, hContext, type, query);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphProfilingPoolCreate(ze_graph_handle_t hGraph,
                           uint32_t count,
                           ze_graph_profiling_pool_handle_t *phProfilingPool) {
    trace_zeGraphProfilingPoolCreate(hGraph, count, phProfilingPool);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret,
                        L0::Graph::fromHandle(hGraph)->createProfilingPool(count, phProfilingPool));

exit:
    trace_zeGraphProfilingPoolCreate(ret, hGraph, count, phProfilingPool);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphProfilingPoolDestroy(ze_graph_profiling_pool_handle_t hProfilingPool) {
    trace_zeGraphProfilingPoolDestroy(hProfilingPool);
    ze_result_t ret;

    if (hProfilingPool == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::GraphProfilingPool::fromHandle(hProfilingPool)->destroy());

exit:
    trace_zeGraphProfilingPoolDestroy(ret, hProfilingPool);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphProfilingQueryCreate(ze_graph_profiling_pool_handle_t hProfilingPool,
                            uint32_t index,
                            ze_graph_profiling_query_handle_t *phProfilingQuery) {
    trace_zeGraphProfilingQueryCreate(hProfilingPool, index, phProfilingQuery);
    ze_result_t ret;

    if (hProfilingPool == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret,
                        L0::GraphProfilingPool::fromHandle(hProfilingPool)
                            ->createProfilingQuery(index, phProfilingQuery));

exit:
    trace_zeGraphProfilingQueryCreate(ret, hProfilingPool, index, phProfilingQuery);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphProfilingQueryDestroy(ze_graph_profiling_query_handle_t hProfilingQuery) {
    trace_zeGraphProfilingQueryDestroy(hProfilingQuery);
    ze_result_t ret;

    if (hProfilingQuery == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::GraphProfilingQuery::fromHandle(hProfilingQuery)->destroy());

exit:
    trace_zeGraphProfilingQueryDestroy(ret, hProfilingQuery);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphProfilingQueryGetData(ze_graph_profiling_query_handle_t hProfilingQuery,
                             ze_graph_profiling_type_t type,
                             uint32_t *pSize,
                             uint8_t *pData) {
    trace_zeGraphProfilingQueryGetData(hProfilingQuery, type, pSize, pData);
    ze_result_t ret;

    if (hProfilingQuery == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(
        ret,
        L0::GraphProfilingQuery::fromHandle(hProfilingQuery)->getData(type, pSize, pData));

exit:
    trace_zeGraphProfilingQueryGetData(ret, hProfilingQuery, type, pSize, pData);
    return ret;
}

ze_result_t ZE_APICALL zeDeviceGetProfilingDataProperties(
    ze_device_handle_t hDevice,
    ze_device_profiling_data_properties_t *pDeviceProfilingDataProperties) {
    trace_zeDeviceGetProfilingDataProperties(hDevice, pDeviceProfilingDataProperties);
    ze_result_t ret;

    if (hDevice == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }
    if (pDeviceProfilingDataProperties == nullptr) {
        LOG_E("Invalid device profiling data properties pointer");
        ret = ZE_RESULT_ERROR_INVALID_NULL_POINTER;
        goto exit;
    }

    if (Driver::getInstance()->getEnvVariables().extensionValidation) {
        if (pDeviceProfilingDataProperties->stype !=
            ZE_STRUCTURE_TYPE_DEVICE_PROFILING_DATA_PROPERTIES) {
            LOG_E("Invalid device profiling data properties stype");
            ret = ZE_RESULT_ERROR_INVALID_ARGUMENT;
            goto exit;
        }
        if (pDeviceProfilingDataProperties->pNext != nullptr) {
            LOG_E("pNext in device profiling data properties not supported");
            ret = ZE_RESULT_ERROR_INVALID_ARGUMENT;
            goto exit;
        }
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::getProfilingDataProperties(pDeviceProfilingDataProperties));

exit:
    trace_zeDeviceGetProfilingDataProperties(ret, hDevice, pDeviceProfilingDataProperties);
    return ret;
}

ze_result_t ZE_APICALL
zeGraphProfilingLogGetString(ze_graph_profiling_query_handle_t hProfilingQuery,
                             uint32_t *pSize,
                             char *pProfilingLog) {
    trace_zeGraphProfilingLogGetString(hProfilingQuery, pSize, pProfilingLog);
    ze_result_t ret;

    if (hProfilingQuery == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(
        ret,
        L0::GraphProfilingQuery::fromHandle(hProfilingQuery)->getLogString(pSize, pProfilingLog));

exit:
    trace_zeGraphProfilingLogGetString(ret, hProfilingQuery, pSize, pProfilingLog);
    return ret;
}

ze_result_t ZE_APICALL zeGraphCompilerGetSupportedOptions(ze_device_handle_t hDevice,
                                                          ze_npu_options_type_t type,
                                                          size_t *pSize,
                                                          char *pSupportedOptions) {
    trace_zeGraphCompilerGetSupportedOptions(type, pSize, pSupportedOptions);
    ze_result_t ret = ZE_RESULT_SUCCESS;
    if (hDevice == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    if (pSize == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_POINTER;
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, Graph::getSupportedOptions(hDevice, type, pSize, pSupportedOptions));

exit:
    trace_zeGraphCompilerGetSupportedOptions(ret, type, pSize, pSupportedOptions);
    return ret;
}

ze_result_t ZE_APICALL zeGraphCompilerIsOptionSupported(ze_device_handle_t hDevice,
                                                        ze_npu_options_type_t type,
                                                        const char *pOption,
                                                        const char *pValue) {
    trace_zeGraphCompilerIsOptionSupported(type, pOption, pValue);
    ze_result_t ret = ZE_RESULT_SUCCESS;
    if (hDevice == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    if (pOption == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_POINTER;
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, Graph::isOptionSupported(hDevice, type, pOption, pValue));

exit:
    trace_zeGraphCompilerIsOptionSupported(ret, type, pOption, pValue);
    return ret;
}

ze_result_t ZE_APICALL zeGraphCreate3(ze_context_handle_t hContext,
                                      ze_device_handle_t hDevice,
                                      const ze_graph_desc_2_t *desc,
                                      ze_graph_handle_t *phGraph,
                                      ze_graph_build_log_handle_t *phGraphBuildLog) {
    trace_zeGraphCreate3(hContext, hDevice, desc, phGraph, phGraphBuildLog);
    ze_result_t ret;

    if (hDevice == nullptr || hContext == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(desc, {ZE_STRUCTURE_TYPE_GRAPH_INPUT_HASH});
    ;
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_CONTEXT, hContext);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    ret = translateHandle(ZEL_HANDLE_DEVICE, hDevice);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::create(hContext, hDevice, desc, phGraph, phGraphBuildLog));

exit:
    trace_zeGraphCreate3(ret, hContext, hDevice, desc, phGraph, phGraphBuildLog);
    return ret;
}

ze_result_t ZE_APICALL zeGraphBuildLogGetString2(ze_graph_build_log_handle_t hGraphBuildLog,
                                                 uint32_t *pSize,
                                                 char *pBuildLog) {
    trace_zeGraphBuildLogGetString2(hGraphBuildLog, pSize, pBuildLog);
    ze_result_t ret;

    if (hGraphBuildLog == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(
        ret,
        L0::GraphBuildLog::fromHandle(hGraphBuildLog)->getLogString(pSize, pBuildLog));

exit:
    trace_zeGraphBuildLogGetString2(ret, hGraphBuildLog, pSize, pBuildLog);
    return ret;
}

ze_result_t ZE_APICALL zeGraphBuildLogDestroy(ze_graph_build_log_handle_t hGraphBuildLog) {
    trace_zeGraphBuildLogDestroy(hGraphBuildLog);
    ze_result_t ret;

    if (hGraphBuildLog == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::GraphBuildLog::fromHandle(hGraphBuildLog)->destroy());

exit:
    trace_zeGraphBuildLogDestroy(ret, hGraphBuildLog);
    return ret;
}

ze_result_t ZE_APICALL zeGraphGetProperties3(ze_graph_handle_t hGraph,
                                             ze_graph_properties_3_t *pGraphProperties) {
    trace_zeGraphGetProperties3(hGraph, pGraphProperties);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    ret = validateGraphStructure(pGraphProperties);
    if (ret != ZE_RESULT_SUCCESS) {
        goto exit;
    }

    L0_HANDLE_EXCEPTION(ret, L0::Graph::fromHandle(hGraph)->getProperties3(pGraphProperties));

exit:
    trace_zeGraphGetProperties3(ret, hGraph, pGraphProperties);
    return ret;
}

ze_result_t ZE_APICALL zeGraphEvict(ze_graph_handle_t hGraph) {
    trace_zeGraphEvict(hGraph);
    ze_result_t ret;

    if (hGraph == nullptr) {
        ret = ZE_RESULT_ERROR_INVALID_NULL_HANDLE;
        goto exit;
    }

    // NOP operation due to lack of kernel module driver support for memory eviction
    ret = ZE_RESULT_SUCCESS;
exit:
    trace_zeGraphEvict(ret, hGraph);
    return ret;
}
} // namespace L0
