/*
 * Copyright (C) 2024-2025 Intel Corporation
 *
 * SPDX-License-Identifier: MIT
 *
 */

#include "shared/source/utilities/staging_buffer_manager.h"

#include "shared/source/command_stream/command_stream_receiver.h"
#include "shared/source/debug_settings/debug_settings_manager.h"
#include "shared/source/device/device.h"
#include "shared/source/execution_environment/root_device_environment.h"
#include "shared/source/helpers/aligned_memory.h"
#include "shared/source/helpers/hw_info.h"
#include "shared/source/memory_manager/unified_memory_manager.h"
#include "shared/source/os_interface/os_interface.h"
#include "shared/source/utilities/heap_allocator.h"

namespace NEO {

StagingBuffer::StagingBuffer(void *baseAddress, size_t size) : baseAddress(baseAddress) {
    this->allocator = std::make_unique<HeapAllocator>(castToUint64(baseAddress), size, MemoryConstants::pageSize, 0u);
}

StagingBuffer::StagingBuffer(StagingBuffer &&other) : baseAddress(other.baseAddress) {
    this->allocator.reset(other.allocator.release());
}

bool StagingBufferTracker::isReady() const {
    return csr->testTaskCountReady(csr->getTagAddress(), taskCountToWait);
}

void StagingBufferTracker::freeChunk() const {
    allocator->free(chunkAddress, size);
}

StagingBufferManager::StagingBufferManager(SVMAllocsManager *svmAllocsManager, const RootDeviceIndicesContainer &rootDeviceIndices, const std::map<uint32_t, DeviceBitfield> &deviceBitfields, bool requiresWritable)
    : svmAllocsManager(svmAllocsManager), rootDeviceIndices(rootDeviceIndices), deviceBitfields(deviceBitfields), requiresWritable(requiresWritable) {
    if (debugManager.flags.StagingBufferSize.get() != -1) {
        chunkSize = debugManager.flags.StagingBufferSize.get() * MemoryConstants::kiloByte;
    }
}

StagingBufferManager::~StagingBufferManager() {
    for (auto &stagingBuffer : stagingBuffers) {
        svmAllocsManager->freeSVMAlloc(stagingBuffer.getBaseAddress());
    }
}

/*
 * This method performs single chunk transfer. If transfer is a read operation, it will fetch oldest staging
 * buffer from the queue, otherwise it allocates or reuses buffer from the pool.
 * After transfer is submitted to GPU, it stores used buffer to either queue in case of reads,
 * or tracking container for further reusage.
 */
template <class Func, class... Args>
StagingTransferStatus StagingBufferManager::performChunkTransfer(size_t chunkTransferId, bool isRead, const UserData &userData, StagingQueue &currentStagingBuffers, CommandStreamReceiver *csr, Func &func, Args... args) {
    StagingTransferStatus result{};
    StagingBufferTracker tracker{};
    auto stagingBufferIndex = chunkTransferId % maxInFlightReads;
    if (isRead && chunkTransferId >= maxInFlightReads) {
        if (copyStagingToHost(currentStagingBuffers[stagingBufferIndex], tracker) == WaitStatus::gpuHang) {
            result.waitStatus = WaitStatus::gpuHang;
            return result;
        }
    } else {
        auto allocatedSize = userData.size;
        auto [allocator, stagingBuffer] = requestStagingBuffer(allocatedSize);
        tracker = StagingBufferTracker{allocator, stagingBuffer, allocatedSize, csr};
    }

    auto stagingBuffer = addrToPtr(tracker.chunkAddress);
    if (!isRead) {
        memcpy(stagingBuffer, userData.ptr, userData.size);
    }

    result.chunkCopyStatus = func(stagingBuffer, args...);

    tracker.taskCountToWait = csr->peekTaskCount();
    if (isRead) {
        currentStagingBuffers[stagingBufferIndex] = {userData, tracker};
    } else {
        trackChunk(tracker);
    }

    if (csr->isAnyDirectSubmissionEnabled()) {
        csr->flushTagUpdate();
    }
    return result;
}

/*
 * This method copies data between non-USM and USM allocations by splitting transfers into chunks.
 * Each chunk copy contains staging buffer which should be used instead of non-usm memory during transfers on GPU.
 * Caller provides actual function to transfer data for single chunk.
 */
StagingTransferStatus StagingBufferManager::performCopy(void *dstPtr, const void *srcPtr, size_t size, ChunkCopyFunction &chunkCopyFunc, CommandStreamReceiver *csr) {
    StagingQueue stagingQueue;
    auto copiesNum = size / chunkSize;
    auto remainder = size % chunkSize;
    StagingTransferStatus result{};
    for (auto i = 0u; i < copiesNum; i++) {
        auto chunkDst = ptrOffset(dstPtr, i * chunkSize);
        auto chunkSrc = ptrOffset(srcPtr, i * chunkSize);
        UserData userData{chunkSrc, chunkSize};
        result = performChunkTransfer(i, false, userData, stagingQueue, csr, chunkCopyFunc, chunkDst, chunkSize);
        if (result.chunkCopyStatus != 0) {
            return result;
        }
    }

    if (remainder != 0) {
        auto chunkDst = ptrOffset(dstPtr, copiesNum * chunkSize);
        auto chunkSrc = ptrOffset(srcPtr, copiesNum * chunkSize);
        UserData userData{chunkSrc, remainder};
        auto result = performChunkTransfer(copiesNum, false, userData, stagingQueue, csr, chunkCopyFunc, chunkDst, remainder);
        if (result.chunkCopyStatus != 0) {
            return result;
        }
    }
    return result;
}

/*
 * This method orchestrates transfer operation for images with given origin and region.
 * Transfer is splitted into chunks, each chunk represents sub-region to transfer.
 * Each chunk contains staging buffer which should be used instead of non-usm memory during transfers on GPU.
 * Several rows are packed into single chunk unless size of single row exceeds maximum chunk size (2MB).
 * Caller provides actual function to enqueue read/write operation for single chunk.
 */
StagingTransferStatus StagingBufferManager::performImageTransfer(const void *ptr, const size_t *globalOrigin, const size_t *globalRegion, size_t rowPitch, size_t bytesPerPixel, ChunkTransferImageFunc &chunkTransferImageFunc, CommandStreamReceiver *csr, bool isRead) {
    StagingQueue stagingQueue;
    size_t origin[3] = {};
    size_t region[3] = {};
    origin[0] = globalOrigin[0];
    origin[2] = globalOrigin[2];
    region[0] = globalRegion[0];
    region[2] = globalRegion[2];
    auto rowsPerChunk = std::max<size_t>(1ul, chunkSize / rowPitch);
    rowsPerChunk = std::min<size_t>(rowsPerChunk, globalRegion[1]);
    auto numOfChunks = globalRegion[1] / rowsPerChunk;
    auto remainder = globalRegion[1] % (rowsPerChunk * numOfChunks);
    StagingTransferStatus result{};
    RowPitchData rowPitchData{region[0] * bytesPerPixel, rowPitch, rowsPerChunk};

    for (auto i = 0u; i < numOfChunks; i++) {
        origin[1] = globalOrigin[1] + i * rowsPerChunk;
        region[1] = rowsPerChunk;
        auto size = region[1] * rowPitch;
        auto chunkPtr = ptrOffset(ptr, i * rowsPerChunk * rowPitch);
        UserData userData{chunkPtr, size, rowPitchData};

        result = performChunkTransfer(i, isRead, userData, stagingQueue, csr, chunkTransferImageFunc, origin, region);
        if (result.chunkCopyStatus != 0 || result.waitStatus == WaitStatus::gpuHang) {
            return result;
        }
    }

    if (remainder != 0) {
        origin[1] = globalOrigin[1] + numOfChunks * rowsPerChunk;
        region[1] = remainder;
        auto size = region[1] * rowPitch;
        auto chunkPtr = ptrOffset(ptr, numOfChunks * rowsPerChunk * rowPitch);
        rowPitchData.rowsInChunk = remainder;
        UserData userData{chunkPtr, size, rowPitchData};

        result = performChunkTransfer(numOfChunks, isRead, userData, stagingQueue, csr, chunkTransferImageFunc, origin, region);
        if (result.chunkCopyStatus != 0 || result.waitStatus == WaitStatus::gpuHang) {
            return result;
        }
    }

    if (isRead) {
        auto numOfSubmittedTransfers = numOfChunks + (remainder != 0 ? 1 : 0);
        result.waitStatus = drainAndReleaseStagingQueue(stagingQueue, std::min(numOfSubmittedTransfers, maxInFlightReads));
    }
    return result;
}

StagingTransferStatus StagingBufferManager::performBufferTransfer(const void *ptr, size_t globalOffset, size_t globalSize, ChunkTransferBufferFunc &chunkTransferBufferFunc, CommandStreamReceiver *csr, bool isRead) {
    StagingQueue stagingQueue;
    auto copiesNum = globalSize / chunkSize;
    auto remainder = globalSize % chunkSize;
    auto chunkOffset = globalOffset;
    StagingTransferStatus result{};
    for (auto i = 0u; i < copiesNum; i++) {
        auto chunkPtr = ptrOffset(ptr, i * chunkSize);
        UserData userData{chunkPtr, chunkSize};
        result = performChunkTransfer(i, isRead, userData, stagingQueue, csr, chunkTransferBufferFunc, chunkOffset, chunkSize);
        if (result.chunkCopyStatus != 0) {
            return result;
        }
        chunkOffset += chunkSize;
    }

    if (remainder != 0) {
        auto chunkPtr = ptrOffset(ptr, copiesNum * chunkSize);
        UserData userData{chunkPtr, remainder};
        result = performChunkTransfer(copiesNum, isRead, userData, stagingQueue, csr, chunkTransferBufferFunc, chunkOffset, remainder);
        if (result.chunkCopyStatus != 0) {
            return result;
        }
    }

    return result;
}

/*
 * This method is used for read transfers. It waits for transfer to finish
 * and copies data associated with that transfer to host allocation.
 * Returned tracker contains staging buffer ready for reuse.
 */
WaitStatus StagingBufferManager::copyStagingToHost(const std::pair<UserData, StagingBufferTracker> &transfer, StagingBufferTracker &tracker) const {
    auto status = transfer.second.csr->waitForTaskCount(transfer.second.taskCountToWait);
    if (status == WaitStatus::gpuHang) {
        return status;
    }
    transfer.second.csr->downloadAllocations(true);
    auto &userData = transfer.first;
    tracker = transfer.second;
    auto stagingBuffer = addrToPtr(tracker.chunkAddress);
    auto userDst = const_cast<void *>(userData.ptr);
    if (userData.rowPitchData.rowSize < userData.rowPitchData.rowPitch) {
        for (auto rowId = 0u; rowId < userData.rowPitchData.rowsInChunk; rowId++) {
            auto offset = rowId * userData.rowPitchData.rowPitch;
            memcpy(ptrOffset(userDst, offset), ptrOffset(stagingBuffer, offset), userData.rowPitchData.rowSize);
        }
    } else {
        memcpy(userDst, stagingBuffer, userData.size);
    }
    return WaitStatus::ready;
}

/*
 * Waits for all pending transfers to finish.
 * Releases staging buffers back to pool for reuse.
 */
WaitStatus StagingBufferManager::drainAndReleaseStagingQueue(const StagingQueue &stagingQueue, size_t numOfTransfers) const {
    StagingBufferTracker tracker{};
    for (auto i = 0u; i < numOfTransfers; i++) {
        auto status = copyStagingToHost(stagingQueue[i], tracker);
        if (status == WaitStatus::gpuHang) {
            return status;
        }
        tracker.freeChunk();
    }
    return WaitStatus::ready;
}

/*
 * This method returns allocator and chunk from staging buffer.
 * Creates new staging buffer if it failed to allocate chunk from existing buffers.
 */
std::pair<HeapAllocator *, uint64_t> StagingBufferManager::requestStagingBuffer(size_t &size) {
    auto lock = std::lock_guard<std::mutex>(mtx);

    auto [allocator, chunkBuffer] = getExistingBuffer(size);
    if (chunkBuffer != 0) {
        return {allocator, chunkBuffer};
    }

    clearTrackedChunks();

    auto [retriedAllocator, retriedChunkBuffer] = getExistingBuffer(size);
    if (retriedChunkBuffer != 0) {
        return {retriedAllocator, retriedChunkBuffer};
    }

    auto stagingBufferSize = alignUp(std::max(chunkSize, size), MemoryConstants::pageSize2M);
    auto usmHost = allocateStagingBuffer(stagingBufferSize);
    if (usmHost != nullptr) {
        StagingBuffer stagingBuffer{usmHost, stagingBufferSize};
        allocator = stagingBuffer.getAllocator();
        chunkBuffer = allocator->allocate(size);
        stagingBuffers.push_back(std::move(stagingBuffer));
    }
    return {allocator, chunkBuffer};
}

/*
 * This method will try to allocate chunk from existing staging buffer.
 * Returns allocator and chunk from consumed staging buffer.
 */
std::pair<HeapAllocator *, uint64_t> StagingBufferManager::getExistingBuffer(size_t &size) {
    uint64_t buffer = 0;
    HeapAllocator *allocator = nullptr;

    for (auto &stagingBuffer : stagingBuffers) {
        allocator = stagingBuffer.getAllocator();
        buffer = allocator->allocate(size);
        if (buffer != 0) {
            if (requiresWritable) {
                auto alloc = svmAllocsManager->getSVMAlloc(stagingBuffer.getBaseAddress())->gpuAllocations.getDefaultGraphicsAllocation();
                alloc->setTbxWritable(true, std::numeric_limits<uint32_t>::max());
                alloc->setAubWritable(true, std::numeric_limits<uint32_t>::max());
            }
            break;
        }
    }
    return {allocator, buffer};
}

void *StagingBufferManager::allocateStagingBuffer(size_t size) {
    SVMAllocsManager::UnifiedMemoryProperties unifiedMemoryProperties(InternalMemoryType::hostUnifiedMemory, 0u, rootDeviceIndices, deviceBitfields);
    auto hostPtr = svmAllocsManager->createHostUnifiedMemoryAllocation(size, unifiedMemoryProperties);
    return hostPtr;
}

bool StagingBufferManager::isValidForCopy(const Device &device, void *dstPtr, const void *srcPtr, size_t size, bool hasDependencies, uint32_t osContextId) {
    auto usmDstData = svmAllocsManager->getSVMAlloc(dstPtr);
    auto usmSrcData = svmAllocsManager->getSVMAlloc(srcPtr);
    bool hostToUsmCopy = usmSrcData == nullptr && usmDstData != nullptr;
    bool isUsedByOsContext = false;
    if (usmDstData) {
        isUsedByOsContext = usmDstData->gpuAllocations.getGraphicsAllocation(device.getRootDeviceIndex())->isUsedByOsContext(osContextId);
    }
    return this->isValidForStaging(device, srcPtr, size, hasDependencies) && hostToUsmCopy && (isUsedByOsContext || size <= chunkSize);
}

bool StagingBufferManager::isValidForStagingTransfer(const Device &device, const void *ptr, size_t size, bool hasDependencies) {
    auto nonUsmPtr = ptr != nullptr && svmAllocsManager->getSVMAlloc(ptr) == nullptr;
    return this->isValidForStaging(device, ptr, size, hasDependencies) && nonUsmPtr;
}

// Common checks for usm, buffers and images
bool StagingBufferManager::isValidForStaging(const Device &device, const void *ptr, size_t size, bool hasDependencies) {
    auto stagingCopyEnabled = device.getProductHelper().isStagingBuffersEnabled();
    if (debugManager.flags.EnableCopyWithStagingBuffers.get() != -1) {
        stagingCopyEnabled = debugManager.flags.EnableCopyWithStagingBuffers.get();
    }
    auto isIntegrated = device.getRootDeviceEnvironment().getHardwareInfo()->capabilityTable.isIntegratedDevice;
    auto osInterface = device.getRootDeviceEnvironment().osInterface.get();
    bool sizeWithinThreshold = osInterface ? osInterface->isSizeWithinThresholdForStaging(size, isIntegrated) : true;
    auto detectedHostPtr = this->registerHostPtr(ptr);
    return stagingCopyEnabled && !hasDependencies && !detectedHostPtr && sizeWithinThreshold;
}

void StagingBufferManager::clearTrackedChunks() {
    for (auto iterator = trackers.begin(); iterator != trackers.end();) {
        if (iterator->isReady()) {
            iterator->freeChunk();
            iterator = trackers.erase(iterator);
        } else {
            break;
        }
    }
}

void StagingBufferManager::trackChunk(const StagingBufferTracker &tracker) {
    auto lock = std::lock_guard<std::mutex>(mtx);
    trackers.push_back(tracker);
}

bool StagingBufferManager::registerHostPtr(const void *ptr) {
    auto lock = std::lock_guard<std::mutex>(mtx);
    auto isHostPtrDetected = detectedHostPtrs.find(ptr) != detectedHostPtrs.end();
    detectedHostPtrs.insert(ptr);
    return isHostPtrDetected;
}

void StagingBufferManager::resetDetectedPtrs() {
    auto lock = std::lock_guard<std::mutex>(mtx);
    detectedHostPtrs.clear();
}

} // namespace NEO
