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

#include "shared/source/debug_settings/debug_settings_manager.h"
#include "shared/source/device/device.h"
#include "shared/source/execution_environment/execution_environment.h"
#include "shared/source/execution_environment/root_device_environment.h"
#include "shared/source/helpers/driver_model_type.h"
#include "shared/source/helpers/hw_info.h"
#include "shared/source/memory_manager/memory_manager.h"
#include "shared/source/memory_manager/unified_memory_manager.h"
#include "shared/source/os_interface/linux/drm_memory_manager.h"
#include "shared/source/os_interface/linux/drm_neo.h"
#include "shared/source/os_interface/linux/ioctl_helper.h"
#include "shared/source/os_interface/linux/ipc_socket_client.h"
#include "shared/source/os_interface/linux/ipc_socket_server.h"
#include "shared/source/os_interface/linux/sys_calls.h"
#include "shared/source/os_interface/os_interface.h"
#include "shared/source/utilities/cpuintrinsics.h"

#include "level_zero/core/source/context/context.h"
#include "level_zero/core/source/device/device.h"
#include "level_zero/core/source/driver/driver_handle.h"

#include <sys/mman.h>
#include <sys/prctl.h>
#include <sys/resource.h>

namespace L0 {

uint8_t Context::isOpaqueHandleSupported(IpcHandleType *handleType) {
    auto rootEnv = this->driverHandle->getMemoryManager()->peekExecutionEnvironment().rootDeviceEnvironments[0].get();
    auto driverType = rootEnv->osInterface ? rootEnv->osInterface->getDriverModel()->getDriverModelType()
                                           : NEO::DriverModelType::unknown;
    if (driverType == NEO::DriverModelType::wddm) {
        *handleType = IpcHandleType::ntHandle;
        return OpaqueHandlingType::nthandle;
    } else {
        *handleType = IpcHandleType::fdHandle;
        for (auto &device : this->driverHandle->devices) {
            auto &productHelper = device->getNEODevice()->getProductHelper();
            if (!productHelper.isPidFdOrSocketForIpcSupported()) {
                return OpaqueHandlingType::none;
            }
        }
        if (NEO::debugManager.flags.ForceIpcSocketFallback.get()) {
            PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                         "Forcing IPC socket method instead of pidfd\n");
            return OpaqueHandlingType::sockets;
        }
        if (NEO::SysCalls::prctl(PR_SET_PTRACER, PR_SET_PTRACER_ANY) == -1) {
            PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr, "prctl Syscall for PR_SET_PTRACER, PR_SET_PTRACER_ANY failed: %s\n", strerror(errno));
            if (NEO::debugManager.flags.EnableIpcSocketFallback.get()) {
                PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                             "pidfd unavailable, but socket fallback is enabled - opaque handles still supported\n");
                return OpaqueHandlingType::sockets;
            }
            return OpaqueHandlingType::none;
        } else {
            // Verify pidfdopen and pidfdgetfd syscalls are available
            int testPidfd = NEO::SysCalls::pidfdopen(-1, 0);
            if (testPidfd == -1 && errno == ENOSYS) {
                PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                             "pidfd_open syscall not available: %s\n", strerror(errno));
                if (NEO::debugManager.flags.EnableIpcSocketFallback.get()) {
                    PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                                 "pidfd unavailable, but socket fallback is enabled - opaque handles still supported\n");
                    return OpaqueHandlingType::sockets;
                }
                return OpaqueHandlingType::none;
            }
            PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                         "pidfd approach for IPC handles is supported\n");
            initOpaqueHandleResources();
            return OpaqueHandlingType::pidfd;
        }
    }
}

bool Context::isIPCHandleSharingSupported() {
    return true;
}

bool Context::isShareableMemory(const void *exportDesc, bool exportableMemory, NEO::Device *neoDevice, bool shareableWithoutNTHandle) {
    if (exportableMemory) {
        return true;
    }
    if (neoDevice->getRootDeviceEnvironment().osInterface) {
        NEO::DriverModelType driverType = neoDevice->getRootDeviceEnvironment().osInterface->getDriverModel()->getDriverModelType();
        if (!exportDesc && driverType == NEO::DriverModelType::wddm) {
            return true;
        }
    }

    return false;
}

void Context::closeExternalHandle(uint64_t handle) {
    auto neoDevice = L0::Device::fromHandle(this->devices.begin()->second)->getNEODevice();
    NEO::DriverModelType driverType = NEO::DriverModelType::unknown;
    if (neoDevice->getRootDeviceEnvironment().osInterface) {
        driverType = neoDevice->getRootDeviceEnvironment().osInterface->getDriverModel()->getDriverModelType();
    }
    if (driverType == NEO::DriverModelType::drm) {
        NEO::SysCalls::close(static_cast<int>(handle));
    }
}

std::pair<NEO::GraphicsAllocation *, void *> Context::getMemHandlePtr(ze_device_handle_t hDevice,
                                                                      uint64_t handle,
                                                                      NEO::AllocationType allocationType,
                                                                      bool isHostIpcAllocation,
                                                                      unsigned int processId,
                                                                      ze_ipc_memory_flags_t flags,
                                                                      uint64_t cacheID,
                                                                      void *reservedHandleData,
                                                                      bool compressedMemory,
                                                                      bool isOpaqueHandle) {
    L0::Device *device = L0::Device::fromHandle(hDevice);
    auto neoDevice = device->getNEODevice();
    NEO::DriverModelType driverType = NEO::DriverModelType::unknown;
    if (neoDevice->getRootDeviceEnvironment().osInterface) {
        driverType = neoDevice->getRootDeviceEnvironment().osInterface->getDriverModel()->getDriverModelType();
    }
    bool isNTHandle = this->getDriverHandle()->getMemoryManager()->isNTHandle(NEO::toOsHandle(reinterpret_cast<void *>(handle)), device->getNEODevice()->getRootDeviceIndex());
    if (isNTHandle) {
        return this->driverHandle->importNTHandle(hDevice,
                                                  reinterpret_cast<void *>(handle),
                                                  allocationType,
                                                  isHostIpcAllocation,
                                                  processId,
                                                  compressedMemory);
    } else if (driverType == NEO::DriverModelType::wddm) {
        return {nullptr, nullptr};
    } else {
        NEO::SvmAllocationData allocDataInternal(neoDevice->getRootDeviceIndex());
        uint64_t effectiveCacheID = cacheID;
        uint64_t importHandle = handle;

        bool opaqueHandlesAttempted = false;
        if (isOpaqueHandle && settings.useOpaqueHandle) {
            // Use helper to import opaque handle with fallback
            auto importResult = importOpaqueHandleWithFallback(handle, processId, cacheID, reservedHandleData, neoDevice);
            if (!importResult.success) {
                return {nullptr, nullptr};
            }
            importHandle = importResult.importHandle;
            opaqueHandlesAttempted = importResult.opaqueHandlesAttempted;
        }

        NEO::GraphicsAllocation *alloc = nullptr;
        auto result = this->driverHandle->importFdHandle(neoDevice,
                                                         flags,
                                                         importHandle,
                                                         allocationType,
                                                         isHostIpcAllocation,
                                                         nullptr,
                                                         &alloc,
                                                         allocDataInternal,
                                                         compressedMemory);
        if (opaqueHandlesAttempted && !alloc && reservedHandleData) {
            result = importHandleFromReservedHandleData(reservedHandleData, cacheID, neoDevice, flags, allocationType, isHostIpcAllocation, compressedMemory, importHandle, alloc);
        }

        // Store cacheID in IPC handle tracking if opaque handles are used
        if (result && isOpaqueHandle && settings.useOpaqueHandle && effectiveCacheID != 0) {
            auto lock = driverHandle->lockIPCHandleMap();
            auto &ipcMap = driverHandle->getIPCHandleMap();
            auto ipcIter = ipcMap.find(importHandle);
            if (ipcIter != ipcMap.end()) {
                ipcIter->second->cacheID = effectiveCacheID;
            }
        }

        return {alloc, result};
    }
}

void Context::getDataFromIpcHandle(ze_device_handle_t hDevice, const ze_ipc_mem_handle_t &ipcHandle, uint64_t &handle, uint8_t &type, unsigned int &processId, uint64_t &poolOffset, uint64_t &cacheID, void *&reservedHandleData, bool &compressedMemory, bool &isOpaqueHandle) {
    if (settings.useOpaqueHandle) {
        const IpcOpaqueMemoryData *ipcData = reinterpret_cast<const IpcOpaqueMemoryData *>(ipcHandle.data);
        handle = static_cast<uint64_t>(ipcData->handle.fd);
        type = ipcData->memoryType;
        processId = ipcData->processId;
        poolOffset = ipcData->poolOffset;
        cacheID = ipcData->computeCacheID();
        compressedMemory = ipcData->compressedMemory;
        uint8_t emptyReservedData[32] = {0};
        if (std::memcmp(ipcData->reservedHandleData, emptyReservedData, sizeof(emptyReservedData)) != 0) {
            reservedHandleData = const_cast<void *>(static_cast<const void *>(&ipcData->reservedHandleData));
        }
        // Check if opaqueHandle matches handle - if not, user changed it, so fallback to legacy path
        uint64_t normalizedHandle = normalizeIPCHandle(*ipcData);
        uint64_t normalizedOpaqueHandle = 0;
        switch (ipcData->type) {
        case IpcHandleType::ntHandle:
            normalizedOpaqueHandle = ipcData->opaqueHandle.reserved;
            break;
        default:
            normalizedOpaqueHandle = static_cast<uint64_t>(static_cast<int64_t>(ipcData->opaqueHandle.fd));
            break;
        }
        isOpaqueHandle = (normalizedHandle == normalizedOpaqueHandle);
    } else {
        const IpcMemoryData *ipcData = reinterpret_cast<const IpcMemoryData *>(ipcHandle.data);
        handle = ipcData->handle;
        type = ipcData->type;
        poolOffset = ipcData->poolOffset;
        processId = 0;
        cacheID = 0;
        reservedHandleData = nullptr;
        isOpaqueHandle = false;
    }
}

ze_result_t Context::systemBarrier(ze_device_handle_t hDevice) {
    if (hDevice == nullptr) {
        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
    }

    auto device = Device::fromHandle(hDevice);
    if (!device) {
        return ZE_RESULT_ERROR_INVALID_ARGUMENT;
    }

    auto neoDevice = device->getNEODevice();

    NEO::DriverModelType driverModelType = NEO::DriverModelType::unknown;
    if (neoDevice->getRootDeviceEnvironment().osInterface) {
        driverModelType = neoDevice->getRootDeviceEnvironment().osInterface->getDriverModel()->getDriverModelType();
    }

    if (driverModelType == NEO::DriverModelType::wddm || driverModelType == NEO::DriverModelType::unknown) {
        return ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
    }

    auto *drmMemoryManager = static_cast<NEO::DrmMemoryManager *>(this->driverHandle->getMemoryManager());

    if (neoDevice->getHardwareInfo().capabilityTable.isIntegratedDevice) {
        return ZE_RESULT_SUCCESS;
    }

    void *ptr = drmMemoryManager->getDrm(neoDevice->getRootDeviceIndex()).getIoctlHelper()->pciBarrierMmap();
    if (ptr == MAP_FAILED || ptr == nullptr) {
        return ZE_RESULT_ERROR_UNKNOWN;
    }

    NEO::CpuIntrinsics::sfence();
    auto pciBarrierPtr = static_cast<uint32_t *>(ptr);
    *pciBarrierPtr = 0u;
    NEO::CpuIntrinsics::sfence();

    NEO::SysCalls::munmap(ptr, MemoryConstants::pageSize);
    return ZE_RESULT_SUCCESS;
}

Context::OpaqueHandleImportResult Context::importOpaqueHandleWithFallback(uint64_t handle,
                                                                          unsigned int processId,
                                                                          uint64_t cacheID,
                                                                          void *reservedHandleData,
                                                                          NEO::Device *neoDevice) {
    uint64_t importHandle = handle;
    bool handleRetrieved = false;
    bool socketFallbackSuccess = false;
    bool opaqueHandlesAttempted = false;

    // Check cache first for opaque handles
    if (this->driverHandle->tryGetCachedImportHandle(cacheID, importHandle)) {
        handleRetrieved = true; // Mark as successful to skip import logic
    }

    // Try pidfd approach first (unless forced to use socket fallback or already cached)
    if (!handleRetrieved && (settings.useOpaqueHandle & OpaqueHandlingType::pidfd) && !NEO::debugManager.flags.ForceIpcSocketFallback.get()) {
        pid_t exporterPid = static_cast<pid_t>(processId);
        unsigned int pidfdFlags = 0u;
        int pidfd = NEO::SysCalls::pidfdopen(exporterPid, pidfdFlags);
        if (pidfd == -1) {
            PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr, "pidfd_open Syscall failed: %s\n", strerror(errno));
        } else {
            unsigned int getfdFlags = 0u;
            int newfd = NEO::SysCalls::pidfdgetfd(pidfd, static_cast<int>(handle), getfdFlags);
            NEO::SysCalls::close(pidfd);
            if (newfd < 0) {
                PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr, "pidfd_getfd Syscall failed: %s\n", strerror(errno));
                return {0, false, false};
            } else {
                importHandle = static_cast<uint64_t>(newfd);
                handleRetrieved = true;
                opaqueHandlesAttempted = true;
                // Cache the imported handle for future use
                this->driverHandle->setCachedImportHandle(cacheID, importHandle);
            }
        }
    }

    // Try socket fallback if pidfd failed and socket fallback is enabled
    if (!handleRetrieved && (settings.useOpaqueHandle & OpaqueHandlingType::sockets)) {
        pid_t exporterPid = static_cast<pid_t>(processId);
        std::string socketPath = "neo_ipc_" + std::to_string(exporterPid);

        NEO::IpcSocketClient socketClient;
        if (socketClient.connectToServer(socketPath)) {
            int receivedFd = socketClient.requestHandle(handle);
            if (receivedFd != -1) {
                importHandle = static_cast<uint64_t>(receivedFd);
                socketFallbackSuccess = true;
                opaqueHandlesAttempted = true;
                // Cache the imported handle for future use
                this->driverHandle->setCachedImportHandle(cacheID, importHandle);
                PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                             "IPC socket fallback successful for handle %lu, cached as %lu\n",
                             handle, importHandle);
            } else {
                PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                             "IPC socket fallback failed for handle %lu\n", handle);
            }
        } else {
            PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                         "Failed to connect to IPC socket server at %s\n", socketPath.c_str());
        }

        // If socket fallback was attempted but failed, return failure
        if (!socketFallbackSuccess) {
            PRINT_STRING(NEO::debugManager.flags.PrintDebugMessages.get(), stderr,
                         "Socket fallback failed for handle %lu, returning failure\n", handle);
            return {0, false, false};
        }
    }

    return {importHandle, true, opaqueHandlesAttempted};
}
} // namespace L0
