#include "iwa/util/texture_atlas.hpp"

#include <bit>
#include "iwa/device.hpp"
#include "iwa/instance.hpp"
#include "iwa/resource/bitmap.hpp"

namespace iwa
{
TextureSlot::TextureSlot(ObjectPtr<TextureAtlas> owner, const TextureSlotCreationArgs& args)
        : super_t(std::move(owner)), mUsedSpace(args.usedSpace), mLayer(args.layer), mUvOffset(args.uvOffset), mUvScale(args.uvScale)
{

}

TextureAtlas::TextureAtlas(ObjectPtr<> owner, const TextureAtlasCreationArgs& args)
    : super_t(std::move(owner)), mLayerSize(args.layerSize)
{
    // start with a single layer with one free space that takes up the entire layer
    mLayers.push_back({
        .freeSpaces = {
            vk::Rect2D{
                .offset = { .x = 0, .y = 0 },
                .extent = args.layerSize
            }
        }
    });
}

ObjectPtr<TextureSlot> TextureAtlas::allocateSlot(vk::Extent2D slotSize)
{
    // only uses multiples of 2
    // TODO: check if this actually improves the results
    const vk::Extent2D size = {
        .width = std::bit_ceil(slotSize.width),
        .height = std::bit_ceil(slotSize.height)
    };

    // check if it can even fit
    if (size.width > mLayerSize.width || size.height > mLayerSize.height) {
        throw std::runtime_error("Cannot allocate texture slot, size too big.");
    }

    // find the best fit (minimize product of "wasted" space)
    unsigned lowestWasteSum = std::numeric_limits<unsigned>::max();
    unsigned lowestWasteProduct = std::numeric_limits<unsigned>::max();
    std::vector<TextureAtlasLayer>::iterator foundLayer = mLayers.end();
    std::vector<vk::Rect2D>::iterator foundSpace;
    for (auto itLayer = mLayers.begin(); itLayer != mLayers.end(); ++itLayer)
    {
        for (auto itSpace = itLayer->freeSpaces.begin(); itSpace != itLayer->freeSpaces.end(); ++itSpace)
        {
            if (itSpace->extent.width < size.width || itSpace->extent.height < size.height) {
                continue;
            }
            const unsigned wasteWidth = itSpace->extent.width - size.width;
            const unsigned wasteHeight = itSpace->extent.height - size.height;
            const unsigned wasteProduct = wasteWidth * wasteHeight;
            if (wasteProduct <= lowestWasteProduct)
            {
                const unsigned wasteSum = wasteWidth + wasteHeight;
                if (wasteProduct < lowestWasteProduct || wasteSum < lowestWasteSum)
                {
                    lowestWasteSum = wasteSum;
                    lowestWasteProduct = wasteProduct;
                    foundLayer = itLayer;
                    foundSpace = itSpace;
                }
            }
        } // for (itLayer->freeSpaces)
    } // for (mLayers)

    // if no space was found, make space
    if (foundLayer == mLayers.end())
    {
        mLayers.resize(mLayers.size() + 1);
        mLayers.back().freeSpaces.push_back({
            .offset = { .x = 0, .y = 0},
            .extent = mLayerSize
        });
        foundLayer = std::prev(mLayers.end());
        foundSpace = foundLayer->freeSpaces.begin();
    }

    // save in case the iterator gets invalidated
    const vk::Rect2D space = *foundSpace;

    // remove it
    foundLayer->freeSpaces.erase(foundSpace);

    // now split the space, if necessary
    const bool splitX = space.extent.width > size.width;
    const bool splitY = space.extent.height > size.height;
    if (splitX)
    {
        foundLayer->freeSpaces.push_back({
            .offset = {
                .x = static_cast<std::int32_t>(space.offset.x + size.width),
                .y = space.offset.y
            },
            .extent = {
                .width = space.extent.width - size.width,
                .height = size.height
            }
        });
    }
    if (splitY)
    {
        foundLayer->freeSpaces.push_back({
            .offset = {
                .x = space.offset.x,
                .y = static_cast<std::int32_t>(space.offset.y + size.height)
            },
            .extent = {
                .width = size.width,
                .height = space.extent.height - size.height
            }
        });
    }
    if (splitX && splitY)
    {
        foundLayer->freeSpaces.push_back({
            .offset = {
                .x = static_cast<std::int32_t>(space.offset.x + size.width),
                .y = static_cast<std::int32_t>(space.offset.y + size.height)
            },
            .extent = {
                .width = space.extent.width - size.width,
                .height = space.extent.height - size.height
            }
        });
    }

    // return the result
    return createChild<TextureSlot>(TextureSlotCreationArgs{
        .usedSpace = {
            .offset = space.offset,
            .extent = slotSize
        },
        .layer = static_cast<unsigned>(std::distance(mLayers.begin(), foundLayer)),
        .uvOffset = {
            static_cast<float>(space.offset.x) / static_cast<float>(mLayerSize.width),
            static_cast<float>(space.offset.y) / static_cast<float>(mLayerSize.height)
        },
        .uvScale = {
            static_cast<float>(slotSize.width) / static_cast<float>(mLayerSize.width),
            static_cast<float>(slotSize.height) / static_cast<float>(mLayerSize.height)
        }
    });
}

AtlasedImage::AtlasedImage(ObjectPtr<Device> owner, const AtlasedImageCreationArgs& args)
    : super_t(std::move(owner)), mAtlas(TextureAtlas::create(TextureAtlasCreationArgs{.layerSize = args.size})),
      mFormat(args.format), mMipLevels(args.mipLevels), mUsage(args.usage | vk::ImageUsageFlagBits::eTransferSrc | vk::ImageUsageFlagBits::eTransferDst)
{
    mImage = allocateImage(args.initialLayers);
    mImageView = mImage->createImageView({
        .viewType = vk::ImageViewType::e2DArray
    });
}

mijin::Task<ObjectPtr<TextureSlot>> AtlasedImage::c_allocateSlot(vk::Extent2D slotSize)
{
    IWA_CORO_ENSURE_MAIN_THREAD(*getOwner()->getOwner());

    ObjectPtr<TextureSlot> slot = mAtlas->allocateSlot(slotSize);
    if (slot->getLayer() >= mImage->getArrayLayers())
    {
        const mijin::TaskMutexLock lock = co_await mImageMutex.c_lock();

        // image is too small, resize it
        // this includes a complete copy of the existing image
        ObjectPtr<Image> newImage = allocateImage(slot->getLayer() + 1);
        ObjectPtr<CommandBuffer> cmdBufferPtr = getOwner()->beginScratchCommandBuffer();
        vk::CommandBuffer cmdBuffer = *cmdBufferPtr;
        mImage->applyTransition(cmdBuffer, IMAGE_TRANSITION_TRANSFER_READ);
        newImage->applyTransition(cmdBuffer, IMAGE_TRANSITION_TRANSFER_WRITE);

        // copy ALL the mip levels
        std::vector<vk::ImageCopy> regions;
        regions.reserve(mImage->getMipLevels());
        for (unsigned level = 0; level < mImage->getMipLevels(); ++level)
        {
            const vk::ImageSubresourceLayers copySubresource{
                .aspectMask     = vk::ImageAspectFlagBits::eColor,
                .mipLevel       = level,
                .baseArrayLayer = 0,
                .layerCount     = mImage->getArrayLayers()
            };
            regions.push_back({
                .srcSubresource = copySubresource,
                .srcOffset = {.x = 0, .y = 0, .z = 0},
                .dstSubresource = copySubresource,
                .dstOffset = {.x = 0, .y = 0, .z = 0},
                .extent = {
                    .width  = mAtlas->getLayerSize().width,
                    .height = mAtlas->getLayerSize().height,
                    .depth  = 1
                }
            });
        }
        cmdBuffer.copyImage(
            /* srcImage       = */ *mImage,
            /* srcImageLayout = */ vk::ImageLayout::eTransferSrcOptimal,
            /* dstImage       = */ *newImage,
            /* dstImageLayout = */ vk::ImageLayout::eTransferDstOptimal,
            /* regions        = */ regions
        );
        co_await getOwner()->endScratchCommandBuffer(cmdBufferPtr);
        mImage = std::move(newImage);
        mImageView = mImage->createImageView({
            .viewType = vk::ImageViewType::e2DArray
        });
        imageRecreated.emit();
    }
    co_return slot;
}

mijin::Task<> AtlasedImage::c_upload(const TextureSlot& slot, const Bitmap& bitmap) const noexcept
{
    IWA_CORO_ENSURE_MAIN_THREAD(*getOwner()->getOwner());
    MIJIN_ASSERT(slot.getUsedSpace().extent.width >= bitmap.getSize().width
                 && slot.getUsedSpace().extent.height >= bitmap.getSize().height, "Can't upload image, invalid size.");

    const mijin::TaskMutexLock lock = co_await mImageMutex.c_lock();
    co_await mImage->c_upload(
        /* bitmap      = */ bitmap,
        /* imageOffset = */ {
            .x = slot.getUsedSpace().offset.x,
            .y = slot.getUsedSpace().offset.y,
            .z = 0
        },
        /* baseLayer = */ slot.getLayer()
    );
}

mijin::Task<> AtlasedImage::c_upload(const TextureSlot& slot, const void* data, std::size_t bytes, const vk::Extent2D& bufferImageSize) const noexcept
{
    IWA_CORO_ENSURE_MAIN_THREAD(*getOwner()->getOwner());
    MIJIN_ASSERT(slot.getUsedSpace().extent.width >= bufferImageSize.width
                 && slot.getUsedSpace().extent.height >= bufferImageSize.height, "Can't upload image, invalid size.");

    const mijin::TaskMutexLock lock = co_await mImageMutex.c_lock();
    co_await mImage->c_upload(
        /* data            = */ data,
        /* bytes           = */ bytes,
        /* bufferImageSize = */ {
            .width = bufferImageSize.width,
            .height = bufferImageSize.height,
            .depth = 1
        },
        /* imageOffset     = */ {
            .x = slot.getUsedSpace().offset.x,
            .y = slot.getUsedSpace().offset.y,
            .z = 0
        },
        /* baseLayer       = */ slot.getLayer()
    );
}

mijin::Task<> AtlasedImage::c_blit(const TextureSlot& slot, Image& srcImage) const noexcept
{
    IWA_CORO_ENSURE_MAIN_THREAD(*getOwner()->getOwner());
    MIJIN_ASSERT(slot.getUsedSpace().extent.width >= srcImage.getSize().width
                 && slot.getUsedSpace().extent.height >= srcImage.getSize().height
                 && srcImage.getSize().depth == 1, "Can't upload image, invalid size.");

    const mijin::TaskMutexLock lock = co_await mImageMutex.c_lock();
    co_await mImage->c_blitFrom(
        /* srcImage    = */ srcImage,
        /* regions     = */ {
            vk::ImageBlit{
                .srcSubresource = DEFAULT_SUBRESOURCE_LAYERS,
                .srcOffsets = std::array{
                    vk::Offset3D{
                        .x = 0, .y = 0, .z = 0
                    },
                    vk::Offset3D{
                        .x = static_cast<std::int32_t>(srcImage.getSize().width),
                        .y = static_cast<std::int32_t>(srcImage.getSize().height),
                        .z = 1
                    }
                },
                .dstSubresource = vk::ImageSubresourceLayers{
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .mipLevel = 0,
                    .baseArrayLayer = slot.getLayer(),
                    .layerCount = 1
                },
                .dstOffsets = std::array{
                    vk::Offset3D{
                        .x = slot.getUsedSpace().offset.x,
                        .y = slot.getUsedSpace().offset.y,
                        .z = 0
                    },
                    vk::Offset3D{
                        .x = slot.getUsedSpace().offset.x + static_cast<std::int32_t>(slot.getUsedSpace().extent.width),
                        .y = slot.getUsedSpace().offset.y + static_cast<std::int32_t>(slot.getUsedSpace().extent.height),
                        .z = 1
                    }
                }
            }
        }
    );
}

mijin::Task<> AtlasedImage::c_blit(const TextureSlot& slot, const Bitmap& bitmap) const noexcept
{
    IWA_CORO_ENSURE_MAIN_THREAD(*getOwner()->getOwner());
    MIJIN_ASSERT(slot.getUsedSpace().extent.width >= bitmap.getSize().width
                 && slot.getUsedSpace().extent.height >= bitmap.getSize().height, "Can't upload image, invalid size.");

    const mijin::TaskMutexLock lock = co_await mImageMutex.c_lock();
    co_await mImage->c_blitFrom(
        /* bitmap      = */ bitmap,
        /* regions     = */ {
            vk::ImageBlit{
                .srcSubresource = DEFAULT_SUBRESOURCE_LAYERS,
                .srcOffsets = std::array{
                    vk::Offset3D{
                        .x = 0, .y = 0, .z = 0
                    },
                    vk::Offset3D{
                        .x = static_cast<std::int32_t>(bitmap.getSize().width),
                        .y = static_cast<std::int32_t>(bitmap.getSize().height),
                        .z = 1
                    }
                },
                .dstSubresource = vk::ImageSubresourceLayers{
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .mipLevel = 0,
                    .baseArrayLayer = slot.getLayer(),
                    .layerCount = 1
                },
                .dstOffsets = std::array{
                    vk::Offset3D{
                        .x = slot.getUsedSpace().offset.x,
                        .y = slot.getUsedSpace().offset.y,
                        .z = 0
                    },
                    vk::Offset3D{
                        .x = slot.getUsedSpace().offset.x + static_cast<std::int32_t>(slot.getUsedSpace().extent.width),
                        .y = slot.getUsedSpace().offset.y + static_cast<std::int32_t>(slot.getUsedSpace().extent.height),
                        .z = 1
                    }
                }
            }
        }
    );
}

mijin::Task<> AtlasedImage::c_copy(const TextureSlot& slot, Image& srcImage) const noexcept
{
    IWA_CORO_ENSURE_MAIN_THREAD(*getOwner()->getOwner());
    MIJIN_ASSERT(slot.getUsedSpace().extent.width >= srcImage.getSize().width
                 && slot.getUsedSpace().extent.height >= srcImage.getSize().height
                 && srcImage.getSize().depth == 1, "Can't upload image, invalid size.");

    const mijin::TaskMutexLock lock = co_await mImageMutex.c_lock();
    co_await mImage->c_copyFrom(
        /* srcImage    = */ srcImage,
        /* regions     = */ {
            vk::ImageCopy{
                .srcSubresource = DEFAULT_SUBRESOURCE_LAYERS,
                .srcOffset = {
                    .x = 0, .y = 0, .z = 0
                },
                .dstSubresource = vk::ImageSubresourceLayers{
                    .aspectMask = vk::ImageAspectFlagBits::eColor,
                    .mipLevel = 0,
                    .baseArrayLayer = slot.getLayer(),
                    .layerCount = 1
                },
                .dstOffset = {
                    .x = slot.getUsedSpace().offset.x,
                    .y = slot.getUsedSpace().offset.y,
                    .z = 0
                },
                .extent = srcImage.getSize()
            }
        }
    );
}

ObjectPtr<Image> AtlasedImage::allocateImage(unsigned layers)
{
    ObjectPtr<Image> image = getOwner()->createChild<Image>(ImageCreationArgs{
        .format = mFormat,
        .extent = {
            .width = mAtlas->getLayerSize().width,
            .height = mAtlas->getLayerSize().height,
            .depth = 1
        },
        .mipLevels = mMipLevels,
        .arrayLayers = layers,
        .usage = mUsage
    });
    image->allocateMemory();
    return image;
}
} // namespace iwa