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

#include "level_zero/core/source/image/image_imp.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/gmm_helper/gmm.h"
#include "shared/source/helpers/basic_math.h"
#include "shared/source/helpers/bindless_heaps_helper.h"
#include "shared/source/helpers/gfx_core_helper.h"
#include "shared/source/helpers/hw_info.h"
#include "shared/source/memory_manager/memory_manager.h"
#include "shared/source/os_interface/product_helper.h"

#include "level_zero/core/source/device/device.h"
#include "level_zero/core/source/driver/driver_handle.h"
#include "level_zero/core/source/image/image_format_desc_helper.h"
#include "level_zero/core/source/image/image_formats.h"

#include "neo_igfxfmid.h"

#include <mutex>

namespace L0 {
ImageAllocatorFn imageFactory[NEO::maxProductEnumValue] = {};

bool isImportedWin32Handle(const ze_image_desc_t *imgDesc) {
    const ze_base_desc_t *extendedDesc = reinterpret_cast<const ze_base_desc_t *>(imgDesc->pNext);
    bool importedWin32Handle = false;
    while (extendedDesc) {
        if (extendedDesc->stype == ZE_STRUCTURE_TYPE_EXTERNAL_MEMORY_IMPORT_WIN32) {
            importedWin32Handle = true;
        }
        extendedDesc = reinterpret_cast<const ze_base_desc_t *>(extendedDesc->pNext);
    }
    return importedWin32Handle;
}

void getImageDescriptorFor3ChEmulation(const ze_image_desc_t *origImgDesc, ze_image_desc_t *imgDesc) {
    *imgDesc = *origImgDesc;
    if (origImgDesc->format.layout == ZE_IMAGE_FORMAT_LAYOUT_16_16_16) {
        imgDesc->format.layout = ZE_IMAGE_FORMAT_LAYOUT_16_16_16_16;
        imgDesc->format.x = ZE_IMAGE_FORMAT_SWIZZLE_R;
        imgDesc->format.y = ZE_IMAGE_FORMAT_SWIZZLE_G;
        imgDesc->format.z = ZE_IMAGE_FORMAT_SWIZZLE_B;
        imgDesc->format.w = ZE_IMAGE_FORMAT_SWIZZLE_1;
    }
    if (origImgDesc->format.layout == ZE_IMAGE_FORMAT_LAYOUT_8_8_8) {
        imgDesc->format.layout = ZE_IMAGE_FORMAT_LAYOUT_8_8_8_8;
        imgDesc->format.x = ZE_IMAGE_FORMAT_SWIZZLE_R;
        imgDesc->format.y = ZE_IMAGE_FORMAT_SWIZZLE_G;
        imgDesc->format.z = ZE_IMAGE_FORMAT_SWIZZLE_B;
        imgDesc->format.w = ZE_IMAGE_FORMAT_SWIZZLE_1;
    }
    return;
}

ImageImp::~ImageImp() {
    if ((isImageView() || imageFromBuffer) && this->device != nullptr) {
        auto rootIndex = this->device->getNEODevice()->getRootDeviceIndex();
        if (bindlessInfo.get() && this->device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[rootIndex]->getBindlessHeapsHelper() != nullptr) {
            this->device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[rootIndex]->getBindlessHeapsHelper()->releaseSSToReusePool(*bindlessInfo);
        }
    }
    if (this->device != nullptr) {
        if (!isImageView() && !imageFromBuffer) {
            this->device->getNEODevice()->getMemoryManager()->freeGraphicsMemory(this->allocation);
        }
        if (implicitArgsAllocation) {
            this->device->getNEODevice()->getMemoryManager()->freeGraphicsMemory(this->implicitArgsAllocation);
        }
    }
}

ze_result_t ImageImp::destroy() {
    if (this->getAllocation() && this->device) {
        auto imageAllocPtr = reinterpret_cast<const void *>(this->getAllocation()->getGpuAddress());
        DriverHandle *driverHandle = this->device->getDriverHandle();

        for (auto peerDevice : driverHandle->devices) {
            this->destroyPeerImages(imageAllocPtr, peerDevice);
        }
    }

    delete this;
    return ZE_RESULT_SUCCESS;
}

ze_result_t ImageImp::destroyPeerImages(const void *ptr, Device *device) {

    std::unique_lock<NEO::SpinLock> lock(device->peerImageAllocationsMutex);

    if (device->peerImageAllocations.find(ptr) != device->peerImageAllocations.end()) {
        delete device->peerImageAllocations[ptr];
        device->peerImageAllocations.erase(ptr);
    }

    return ZE_RESULT_SUCCESS;
}

ze_result_t ImageImp::createView(Device *device, const ze_image_desc_t *desc, ze_image_handle_t *pImage) {
    auto productFamily = device->getNEODevice()->getHardwareInfo().platform.eProductFamily;

    ImageAllocatorFn allocator = nullptr;
    allocator = imageFactory[productFamily];

    ImageImp *image = nullptr;

    image = static_cast<ImageImp *>((*allocator)());
    image->allocation = allocation;
    image->sourceImageFormatDesc = this->imageFormatDesc;
    image->imgInfo = this->imgInfo;
    image->imageFromBuffer = this->imageFromBuffer;

    auto result = ZE_RESULT_SUCCESS;
    switch (desc->format.layout) {
    default:
        result = image->initialize(device, desc);
        break;
    case ZE_IMAGE_FORMAT_LAYOUT_32_32_32:
        result = ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
        break;
    case ZE_IMAGE_FORMAT_LAYOUT_8_8_8:
    case ZE_IMAGE_FORMAT_LAYOUT_16_16_16:
        if (isImportedWin32Handle(desc)) {
            result = ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
        } else {
            ze_image_desc_t imgDesc = {};
            getImageDescriptorFor3ChEmulation(desc, &imgDesc);
            image->setMimickedImage(true);
            result = image->initialize(device, &imgDesc);
        }
        break;
    }
    if (result != ZE_RESULT_SUCCESS) {
        image->destroy();
        image = nullptr;
    }
    *pImage = image;

    return result;
}

ze_result_t ImageImp::allocateBindlessSlot() {
    if (!isImageView() && !imageFromBuffer) {
        if (!this->device->getNEODevice()->getMemoryManager()->allocateBindlessSlot(allocation)) {
            return ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY;
        }
        if (allocation->getBindlessOffset() != std::numeric_limits<uint64_t>::max()) {
            bindlessInfo = std::make_unique<NEO::SurfaceStateInHeapInfo>(allocation->getBindlessInfo());
        }
        return ZE_RESULT_SUCCESS;
    }
    auto bindlessHelper = this->device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[allocation->getRootDeviceIndex()]->getBindlessHeapsHelper();

    if (bindlessHelper && !bindlessInfo) {
        auto &gfxCoreHelper = this->device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[allocation->getRootDeviceIndex()]->getHelper<NEO::GfxCoreHelper>();
        const auto surfStateCount = NEO::BindlessImageSlot::max;
        auto surfaceStateSize = surfStateCount * gfxCoreHelper.getRenderSurfaceStateSize();

        auto surfaceStateInfo = bindlessHelper->allocateSSInHeap(surfaceStateSize, allocation, NEO::BindlessHeapsHelper::globalSsh);
        if (surfaceStateInfo.heapAllocation == nullptr) {
            return ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY;
        }
        bindlessInfo = std::make_unique<NEO::SurfaceStateInHeapInfo>(surfaceStateInfo);
    }
    return ZE_RESULT_SUCCESS;
}

NEO::SurfaceStateInHeapInfo *ImageImp::getBindlessSlot() {
    return bindlessInfo.get();
}

ze_result_t Image::create(uint32_t productFamily, Device *device, const ze_image_desc_t *desc, Image **pImage) {
    ze_result_t result = ZE_RESULT_SUCCESS;
    ImageAllocatorFn allocator = nullptr;
    if (productFamily < NEO::maxProductEnumValue) {
        allocator = imageFactory[productFamily];
    }

    ImageImp *image = nullptr;
    if (allocator) {
        image = static_cast<ImageImp *>((*allocator)());
        switch (desc->format.layout) {
        default:
            result = image->initialize(device, desc);
            break;
        case ZE_IMAGE_FORMAT_LAYOUT_8_8_8:
        case ZE_IMAGE_FORMAT_LAYOUT_16_16_16:
            if (isImportedWin32Handle(desc)) {
                result = ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
            } else {
                ze_image_desc_t imgDesc = {};
                getImageDescriptorFor3ChEmulation(desc, &imgDesc);
                image->setMimickedImage(true);
                result = image->initialize(device, &imgDesc);
            }
            break;
        case ZE_IMAGE_FORMAT_LAYOUT_32_32_32:
            result = ZE_RESULT_ERROR_UNSUPPORTED_FEATURE;
            break;
        }
        if (result != ZE_RESULT_SUCCESS) {
            image->destroy();
            image = nullptr;
        }
    } else {
        result = ZE_RESULT_ERROR_UNKNOWN;
    }
    *pImage = image;

    return result;
}

ze_result_t Image::getPitchFor2dImage(
    ze_device_handle_t hDevice,
    size_t imageWidth,
    size_t imageHeight,
    unsigned int elementSizeInByte,
    size_t *rowPitch) {

    NEO::ImageInfo imgInfo = {};
    imgInfo.imgDesc.imageType = NEO::ImageType::image2D;
    imgInfo.imgDesc.imageWidth = imageWidth;
    imgInfo.imgDesc.imageHeight = imageHeight;
    imgInfo.linearStorage = true;
    [[maybe_unused]] uint32_t exponent;
    switch (elementSizeInByte) {
    default:
        exponent = Math::log2(elementSizeInByte);
        if (exponent >= 5u) {
            return ZE_RESULT_ERROR_INVALID_ARGUMENT;
        }
        imgInfo.surfaceFormat = &ImageFormats::surfaceFormatsForRedescribe[exponent % 5];
        break;
    case 3:
        imgInfo.surfaceFormat = &ImageFormats::surfaceFormatsForRedescribe[5];
        break;
    case 6:
        imgInfo.surfaceFormat = &ImageFormats::surfaceFormatsForRedescribe[6];
        break;
    }

    Device *device = Device::fromHandle(hDevice);
    *rowPitch = ImageImp::getRowPitchFor2dImage(device, imgInfo);

    return ZE_RESULT_SUCCESS;
}

ze_result_t ImageImp::getDeviceOffset(uint64_t *deviceOffset) {
    if (!this->bindlessImage) {
        return ZE_RESULT_ERROR_NOT_AVAILABLE;
    }

    DEBUG_BREAK_IF(this->getBindlessSlot() == nullptr);
    *deviceOffset = this->getBindlessSlot()->surfaceStateOffset;

    return ZE_RESULT_SUCCESS;
}

size_t ImageImp::getRowPitchFor2dImage(Device *device, const NEO::ImageInfo &imgInfo) {
    NEO::StorageInfo storageInfo = {};
    NEO::ImageInfo info = imgInfo;

    NEO::Gmm gmm(device->getNEODevice()->getExecutionEnvironment()->rootDeviceEnvironments[device->getRootDeviceIndex()]->getGmmHelper(),
                 info,
                 storageInfo,
                 false);

    return info.rowPitch;
}

ze_result_t ImageImp::allocateImplicitArgsOnDemand() {
    std::lock_guard<std::mutex> lock(implicitArgsAllocationMutex);
    if (implicitArgsAllocation != nullptr) {
        return ZE_RESULT_SUCCESS;
    }

    auto neoDevice = device->getNEODevice();
    auto memoryManager = neoDevice->getMemoryManager();

    const size_t implicitArgsSize = NEO::ImageImplicitArgs::getSize();
    NEO::AllocationProperties properties{
        neoDevice->getRootDeviceIndex(),
        implicitArgsSize,
        NEO::AllocationType::buffer,
        neoDevice->getDeviceBitfield()};

    implicitArgsAllocation = memoryManager->allocateGraphicsMemoryWithProperties(properties);
    if (implicitArgsAllocation == nullptr) {
        return ZE_RESULT_ERROR_OUT_OF_HOST_MEMORY;
    }

    NEO::ImageImplicitArgs imageImplicitArgs{};
    populateImageImplicitArgs(imageImplicitArgs);

    auto &rootDeviceEnvironment = neoDevice->getRootDeviceEnvironment();
    auto &productHelper = rootDeviceEnvironment.getHelper<NEO::ProductHelper>();

    bool success = NEO::MemoryTransferHelper::transferMemoryToAllocation(
        productHelper.isBlitCopyRequiredForLocalMemory(rootDeviceEnvironment, *implicitArgsAllocation),
        *neoDevice,
        implicitArgsAllocation,
        0,
        &imageImplicitArgs,
        implicitArgsSize);

    if (!success) {
        memoryManager->freeGraphicsMemory(implicitArgsAllocation);
        implicitArgsAllocation = nullptr;
        return ZE_RESULT_ERROR_UNKNOWN;
    }
    encodeImplicitArgsSurfaceState();
    return ZE_RESULT_SUCCESS;
}

void ImageImp::populateImageImplicitArgs(NEO::ImageImplicitArgs &imageImplicitArgs) {
    imageImplicitArgs.structVersion = 0;
    imageImplicitArgs.imageWidth = imgInfo.imgDesc.imageWidth;
    imageImplicitArgs.imageHeight = imgInfo.imgDesc.imageHeight;
    imageImplicitArgs.imageDepth = imgInfo.imgDesc.imageDepth;
    imageImplicitArgs.imageArraySize = imgInfo.imgDesc.imageArraySize;
    imageImplicitArgs.numSamples = imgInfo.imgDesc.numSamples;
    imageImplicitArgs.numMipLevels = imgInfo.imgDesc.numMipLevels;

    imageImplicitArgs.channelType = getClChannelDataType(imageFormatDesc.format);
    imageImplicitArgs.channelOrder = getClChannelOrder(imageFormatDesc.format);

    auto pixelSize = imgInfo.surfaceFormat->imageElementSizeInBytes;
    imageImplicitArgs.flatBaseOffset = implicitArgsAllocation->getGpuAddress();
    imageImplicitArgs.flatWidth = (imgInfo.imgDesc.imageWidth * pixelSize) - 1u;
    imageImplicitArgs.flagHeight = (imgInfo.imgDesc.imageHeight * pixelSize) - 1u;
    imageImplicitArgs.flatPitch = imgInfo.imgDesc.imageRowPitch - 1u;
}
} // namespace L0
