#include "iwa/image.hpp"

#include <cmath>
#include "iwa/resource/bitmap.hpp"
#include "iwa/buffer.hpp"
#include "iwa/device.hpp"
#include "iwa/util/vkutil.hpp"

namespace iwa
{
Image::Image(ObjectPtr<Device> owner, ImageCreationArgs args) : super_t(std::move(owner)), mFlags(args.flags), mType(args.imageType),
        mFormat(args.format), mTiling(args.tiling), mUsage(args.usage), mSize(args.extent), mArrayLayers(args.arrayLayers),
        mMipLevels(clampMipLevels(args.mipLevels))
{
    if (mMipLevels > 1) {
        mUsage |= vk::ImageUsageFlagBits::eTransferDst | vk::ImageUsageFlagBits::eTransferSrc;
    }
    mHandle = getOwner()->getVkHandle().createImage(vk::ImageCreateInfo{
        .flags = args.flags,
        .imageType = args.imageType,
        .format = args.format,
        .extent = args.extent,
        .mipLevels = mMipLevels,
        .arrayLayers = args.arrayLayers,
        .samples = args.samples,
        .tiling = args.tiling,
        .usage = mUsage,
        .sharingMode = args.sharingMode,
        .queueFamilyIndexCount = static_cast<std::uint32_t>(args.queueFamilyIndices.size()),
        .pQueueFamilyIndices = args.queueFamilyIndices.data(),
        .initialLayout = args.initialLayout
    });
}

Image::Image(ObjectPtr<Device> owner, ImageWrapArgs args)
    : super_t(std::move(owner)), MixinVulkanObject(args.handle), mType(args.type), mFormat(args.format), mUsage(args.usage),
      mSize(args.size), mMipLevels(args.mipLevels), mWrapped(true)
{

}

Image::~Image() noexcept
{
    if (!mWrapped)
    {
        IWA_DELETE_DEVICE_OBJECT(getOwner(), mHandle, destroyImage)
    }
}

void Image::resetUsage(ResetLayout resetLayout) noexcept
{
    lastUsageStages = vk::PipelineStageFlagBits::eTopOfPipe;
    lastAccess = vk::AccessFlags();
    if (resetLayout)
    {
        currentLayout = vk::ImageLayout::eUndefined;
    }
}

void Image::allocateMemory()
{
    const vk::MemoryRequirements memoryRequirements = getOwner()->getVkHandle().getImageMemoryRequirements(mHandle);
    const vk::MemoryPropertyFlags memoryFlags = vk::MemoryPropertyFlagBits::eDeviceLocal;
    const std::optional<std::uint32_t> memoryTypeIdx = findMemoryType(*getOwner(), memoryRequirements, memoryFlags);
    if (!memoryTypeIdx.has_value())
    {
        throw std::runtime_error("Could not find a suitable memory type.");
    }

    ObjectPtr<DeviceMemory> memory = getOwner()->allocateDeviceMemory(
    {
        .allocationSize = memoryRequirements.size,
        .memoryTypeIndex = memoryTypeIdx.value()
    });
    bindMemory(std::move(memory));
}

void Image::bindMemory(ObjectPtr<DeviceMemory> memory, vk::DeviceSize offset)
{
    mMemory = std::move(memory);
    getOwner()->getVkHandle().bindImageMemory(mHandle, *mMemory, offset);
}

void Image::applyTransition(vk::CommandBuffer cmdBuffer, const ImageTransition& transition)
{
    assert(transition.layout != vk::ImageLayout::eUndefined);

    if (transition.layout != currentLayout)
    {
        cmdBuffer.pipelineBarrier(
            /* srcStageMask = */ lastUsageStages,
            /* dstStageMask = */ transition.stages,
            /* dependencyFlags = */ vk::DependencyFlags(),
            /* memoryBarriers = */ {},
            /* bufferMemoryBarriers = */ {},
            /* imageMemoryBarriers = */ {
            vk::ImageMemoryBarrier{
                .srcAccessMask = lastAccess,
                .dstAccessMask = transition.access,
                .oldLayout = currentLayout,
                .newLayout = transition.layout,
                .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
                .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
                .image = mHandle,
                .subresourceRange = transition.subResourceRange
            }
        });
    }

    lastUsageStages = transition.stages;
    lastAccess = transition.access;
    currentLayout = transition.layout;
}

ObjectPtr<ImageView> Image::createImageView(const ImageViewCreationArgs& args)
{
    return createChild<ImageView>(args);
}


mijin::Task<> Image::c_upload(const void* data, std::size_t bytes, vk::Extent3D bufferImageSize, vk::Offset3D imageOffset,
                              unsigned baseLayer, unsigned layerCount)
{
    MIJIN_ASSERT(bytes >= static_cast<std::size_t>(bufferImageSize.width) * bufferImageSize.height * bufferImageSize.depth
        * vkFormatSize(mFormat), "Buffer for image upload is too small!");

    // TODO: optimize this whole process
    // create scratch buffer
    vk::Device device = *getOwner();
    ObjectPtr<Buffer> scratchBuffer = getOwner()->createChild<Buffer>(BufferCreationArgs{
        .size = bytes,
        .usage = vk::BufferUsageFlagBits::eTransferSrc
    });
    scratchBuffer->allocateMemory(HostVisible::YES);

    // copy to scratch buffer
    void* mapped = device.mapMemory(*scratchBuffer->getMemory(), 0, bytes);
    std::memcpy(mapped, data, bytes);
    device.unmapMemory(*scratchBuffer->getMemory());

    ObjectPtr<CommandBuffer> cmdBufferPtr = getOwner()->beginScratchCommandBuffer();
    const vk::CommandBuffer cmdBuffer = *cmdBufferPtr;
    applyTransition(cmdBuffer, {
        .stages = vk::PipelineStageFlagBits::eTransfer,
        .layout = vk::ImageLayout::eTransferDstOptimal,
        .access = vk::AccessFlagBits::eTransferWrite
    });

    // copy to actual buffer
    cmdBuffer.copyBufferToImage(*scratchBuffer, mHandle, vk::ImageLayout::eTransferDstOptimal, vk::BufferImageCopy{
        .bufferOffset = 0,
        .bufferRowLength = 0,
        .bufferImageHeight = 0,
        .imageSubresource = {
            .aspectMask = vk::ImageAspectFlagBits::eColor,
            .mipLevel = 0,
            .baseArrayLayer = baseLayer,
            .layerCount = layerCount
        },
        .imageOffset = {
            .x = imageOffset.x,
            .y = imageOffset.y,
            .z = imageOffset.z
        },
        .imageExtent = {
            .width = bufferImageSize.width,
            .height = bufferImageSize.height,
            .depth = bufferImageSize.depth
        }
    });
    if (mMipLevels > 1)
    {
        generateMipMaps(cmdBuffer);
    }
    co_await getOwner()->endScratchCommandBuffer(std::move(cmdBufferPtr));
}

mijin::Task<> Image::c_doTransition(const ImageTransition& transition)
{
    ObjectPtr<CommandBuffer> cmdBufferPtr = getOwner()->beginScratchCommandBuffer();
    const vk::CommandBuffer cmdBuffer = *cmdBufferPtr;
    applyTransition(cmdBuffer, transition);

    co_await getOwner()->endScratchCommandBuffer(std::move(cmdBufferPtr));
}

mijin::Task<> Image::c_upload(const Bitmap& bitmap, vk::Offset3D imageOffset, unsigned baseLayer, unsigned layerCount)
{
    MIJIN_ASSERT(vkFormatSize(mFormat) == vkFormatSize(bitmap.getFormat()), "Bitmap format size doesn't match image format size.");

    const vk::Extent3D bufferImageSize = {
        .width = bitmap.getSize().width,
        .height = bitmap.getSize().height,
        .depth = 1
    };
    return c_upload(bitmap.getData().data(), bitmap.getData().size_bytes(), bufferImageSize, imageOffset, baseLayer, layerCount);
}

mijin::Task<> Image::c_blitFrom(Image& srcImage, std::vector<vk::ImageBlit> regions, vk::Filter filter)
{
    ObjectPtr<CommandBuffer> cmdBufferPtr = getOwner()->beginScratchCommandBuffer();
    const vk::CommandBuffer cmdBuffer = *cmdBufferPtr;

    applyTransition(cmdBuffer, IMAGE_TRANSITION_TRANSFER_WRITE);
    srcImage.applyTransition(cmdBuffer, IMAGE_TRANSITION_TRANSFER_READ);

    cmdBuffer.blitImage(
        /* srcImage       = */ srcImage,
        /* srcImageLayout = */ vk::ImageLayout::eTransferSrcOptimal,
        /* dstImage       = */ *this,
        /* dstImageLayout = */ vk::ImageLayout::eTransferDstOptimal,
        /* regions        = */ regions,
        /* filter         = */ filter
    );
    co_await getOwner()->endScratchCommandBuffer(std::move(cmdBufferPtr));
}

mijin::Task<> Image::c_blitFrom(const Bitmap& bitmap, std::vector<vk::ImageBlit> regions, vk::Filter filter)
{
    ObjectPtr<Image> scratchImage = co_await c_create(getOwner()->getPointer(), ImageFromBitmapArgs{
        .bitmap = &bitmap,
        .usage = vk::ImageUsageFlagBits::eTransferDst | vk::ImageUsageFlagBits::eTransferSrc
    });
    co_await c_blitFrom(*scratchImage, std::move(regions), filter);
}

mijin::Task<> Image::c_copyFrom(Image& srcImage, std::vector<vk::ImageCopy> regions)
{
    ObjectPtr<CommandBuffer> cmdBufferPtr = getOwner()->beginScratchCommandBuffer();
    const vk::CommandBuffer cmdBuffer = *cmdBufferPtr;

    applyTransition(cmdBuffer, IMAGE_TRANSITION_TRANSFER_WRITE);
    srcImage.applyTransition(cmdBuffer, IMAGE_TRANSITION_TRANSFER_READ);

    cmdBuffer.copyImage(
        /* srcImage       = */ srcImage,
        /* srcImageLayout = */ vk::ImageLayout::eTransferSrcOptimal,
        /* dstImage       = */ *this,
        /* dstImageLayout = */ vk::ImageLayout::eTransferDstOptimal,
        /* regions        = */ regions
    );
    co_await getOwner()->endScratchCommandBuffer(std::move(cmdBufferPtr));
}

std::uint32_t Image::clampMipLevels(std::uint32_t levels) const
{
    if (levels <= 1) {
        return 1;
    }
    const vk::ImageFormatProperties props = getOwner()->getVkPhysicalDevice().getImageFormatProperties(mFormat, mType, mTiling, mUsage, mFlags);
    return std::min(levels, std::min(props.maxMipLevels, static_cast<std::uint32_t>(std::log2(std::max(mSize.width, mSize.height)))+1));
}

void Image::generateMipMaps(vk::CommandBuffer cmdBuffer)
{
    std::vector<vk::ImageBlit> regions;
    regions.resize(mMipLevels - 1);

    applyTransition(cmdBuffer, {
        .stages = vk::PipelineStageFlagBits::eTransfer,
        .layout = vk::ImageLayout::eGeneral, // TODO: using transfer dst/src optimal would be better, but I don't want to deal with different layout in the mips
        .access = vk::AccessFlagBits::eTransferWrite | vk::AccessFlagBits::eTransferRead
    });

    unsigned mipWidth = mSize.width / 2;
    unsigned mipHeight = mSize.height / 2;
    for (unsigned level = 1; level < mMipLevels; ++level)
    {
        vk::ImageBlit& region = regions[level - 1];
        region.srcSubresource = {
            .aspectMask = vk::ImageAspectFlagBits::eColor,
            .mipLevel = 0,
            .baseArrayLayer = 0,
            .layerCount = 1
        };
        region.srcOffsets[0].x = 0;
        region.srcOffsets[0].y = 0;
        region.srcOffsets[0].z = 0;
        region.srcOffsets[1].x = static_cast<std::int32_t>(mSize.width);
        region.srcOffsets[1].y = static_cast<std::int32_t>(mSize.height);
        region.srcOffsets[1].z = 1;
        region.dstSubresource = {
            .aspectMask = vk::ImageAspectFlagBits::eColor,
            .mipLevel = level,
            .baseArrayLayer = 0,
            .layerCount = 1
        };
        region.dstOffsets[0].x = 0;
        region.dstOffsets[0].y = 0;
        region.dstOffsets[0].z = 0;
        region.dstOffsets[1].x = static_cast<std::int32_t>(mipWidth);
        region.dstOffsets[1].y = static_cast<std::int32_t>(mipHeight);
        region.dstOffsets[1].z = 1;

        if (mipWidth > 1) {
            mipWidth /= 2;
        }
        if (mipHeight > 1) {
            mipHeight /= 2;
        }
    }

    cmdBuffer.blitImage(
        /* srcImage       = */ mHandle,
        /* srcImageLayout = */ vk::ImageLayout::eGeneral,
        /* dstImage       = */ mHandle,
        /* dstImageLayout = */ vk::ImageLayout::eGeneral,
        /* regions        = */ regions,
        /* filter         = */ vk::Filter::eLinear
    );
}

mijin::Task<ObjectPtr<Image>> Image::c_create(ObjectPtr<Device> owner, ImageFromBitmapArgs args)
{
    ObjectPtr<Image> image = owner->createChild<Image>(ImageCreationArgs{
        .flags = args.flags,
        .imageType = vk::ImageType::e2D,
        .format = args.bitmap->getFormat(),
        .extent = {
            .width = args.bitmap->getSize().width,
            .height = args.bitmap->getSize().height,
            .depth = 1
        },
        .mipLevels = 1,
        .arrayLayers = 1,
        .samples = vk::SampleCountFlagBits::e1,
        .tiling = args.tiling,
        .usage = args.usage | vk::ImageUsageFlagBits::eTransferDst,
        .sharingMode = args.sharingMode,
        .queueFamilyIndices = std::move(args.queueFamilyIndices),
        .initialLayout = args.initialLayout
    });
    image->allocateMemory();
    co_await image->c_upload(*args.bitmap);
    co_return image;
}

ImageView::ImageView(ObjectPtr<Image> owner, const ImageViewCreationArgs& args) : super_t(std::move(owner))
{
    mHandle = getOwner()->getOwner()->getVkHandle().createImageView(vk::ImageViewCreateInfo
    {
        .flags = args.flags,
        .image = *getOwner(),
        .viewType = args.viewType,
        .format = args.format != vk::Format::eUndefined ? args.format : getOwner()->getFormat(),
        .components = args.components,
        .subresourceRange = args.subresourceRange
    });
}

ImageView::~ImageView() noexcept
{
    IWA_DELETE_DEVICE_OBJECT(getOwner()->getOwner(), mHandle, destroyImageView);
}

Sampler::Sampler(ObjectPtr<Device> owner, const SamplerCreationArgs& args) : super_t(std::move(owner))
{
    mHandle = getOwner()->getVkHandle().createSampler(vk::SamplerCreateInfo
    {
        .flags = args.flags,
        .magFilter = args.magFilter,
        .minFilter = args.minFilter,
        .mipmapMode = args.mipmapMode,
        .addressModeU = args.addressModeU,
        .addressModeV = args.addressModeV,
        .addressModeW = args.addressModeW,
        .mipLodBias = args.mipLodBias,
        .anisotropyEnable = args.options.anisotropyEnable,
        .maxAnisotropy = args.maxAnisotropy,
        .compareEnable = args.options.compareEnable,
        .compareOp = args.compareOp,
        .minLod = args.minLod,
        .maxLod = args.maxLod,
        .borderColor = args.borderColor,
        .unnormalizedCoordinates = args.options.unnormalizedCoordinates
    });
}

Sampler::~Sampler() noexcept
{
    IWA_DELETE_DEVICE_OBJECT(getOwner(), mHandle, destroySampler);
}
} // namespace iwa