#include "iwa/util/image_reference.hpp"

#include <mijin/util/iterators.hpp>
#include "iwa/command.hpp"
#include "iwa/device.hpp"

namespace iwa
{
ImageReference::ImageReference(ObjectPtr<Device> owner) : super_t(std::move(owner))
{

}

void ImageReference::finalize(ImageReferenceFinalizeArgs& /* args */) {}

mijin::Task<> ImageReference::c_present()
{
    co_return;
}

SwapchainImageReference::SwapchainImageReference(ObjectPtr<Device> owner, SwapchainImageReferenceCreationArgs args)
        : super_t(std::move(owner)), mSwapchain(std::move(args.swapchain))
{
    mPresentReadySemaphores.resize(mSwapchain->getNumParallelFrames());
    for (ObjectPtr<Semaphore>& semaphore : mPresentReadySemaphores)
    {
        semaphore = getOwner()->createChild<Semaphore>();
    }
    createImageViews();
    mSwapchain->recreated.connect([this]()
    {
        createImageViews();
    });
}

vk::Format SwapchainImageReference::getFormat()
{
    return mSwapchain->getFormat();
}

vk::Extent2D SwapchainImageReference::getExtent()
{
    return mSwapchain->getExtent();
}

ImageReferenceFrame SwapchainImageReference::getCurrentFrame()
{
    return ImageReferenceFrame{
        .image = mSwapchain->getCurrentImage().getRaw(),
        .imageView = mImageViews[mSwapchain->getCurrentImageIdx()].getRaw()
    };
}

void SwapchainImageReference::finalize(ImageReferenceFinalizeArgs& args)
{
    args.waitSemaphores.push_back(*mSwapchain->getCurrentAvailableSemaphore());
    args.signalSemaphores.push_back(*mPresentReadySemaphores[mSwapchain->getCurrentFrameIdx()]);

    mSwapchain->getCurrentImage()->applyTransition(args.cmdBuffer, ImageTransition{
        .stages = vk::PipelineStageFlagBits::eBottomOfPipe,
        .layout = vk::ImageLayout::ePresentSrcKHR,
        .access = {}
    });
}

mijin::Task<> SwapchainImageReference::c_present()
{
    // and present
    co_await mSwapchain->c_present({
        .queue = getOwner()->getGraphicsQueue(),
        .waitSemaphores = {mPresentReadySemaphores[mSwapchain->getCurrentFrameIdx()]->getVkHandle()}
    });
}

void SwapchainImageReference::createImageViews()
{
    mImageViews.resize(mSwapchain->getImages().size());
    for (auto [image, imageView] : mijin::zip(mSwapchain->getImages(), mImageViews))
    {
        imageView = image->createImageView();
    }
}


DirectImageReference::DirectImageReference(ObjectPtr<Device> owner, DirectImageReferenceCreationArgs args)
        : super_t(std::move(owner)), mImage(std::move(args.image)), mImageView(std::move(args.imageView))
{

}

vk::Format DirectImageReference::getFormat()
{
    return mImage->getFormat();
}

vk::Extent2D DirectImageReference::getExtent()
{
    return {
        .width = mImage->getSize().width,
        .height = mImage->getSize().height
    };
}

ImageReferenceFrame DirectImageReference::getCurrentFrame()
{
    return ImageReferenceFrame{
        .image = mImage.getRaw(),
        .imageView = mImageView.getRaw()
    };
}

AutoResizeImageReference::AutoResizeImageReference(ObjectPtr<Device> owner, AutoResizeImageReferenceCreationArgs args)
    : super_t(std::move(owner), DirectImageReferenceCreationArgs{}), mReferenceImageRef(std::move(args.referenceImageRef)),
      mImageCreationArgs(std::move(args.imageCreationArgs)), mImageViewCreationArgs(args.imageViewCreationArgs)
{
    createImage();
}

vk::Extent2D AutoResizeImageReference::getExtent()
{
    return mReferenceImageRef->getExtent();
}

ImageReferenceFrame AutoResizeImageReference::getCurrentFrame()
{
    const vk::Extent2D extent = mReferenceImageRef->getExtent();
    if (extent.width != mImage->getSize().width || extent.height != mImage->getSize().height) {
        createImage();
    }
    return ImageReferenceFrame{
        .image = mImage.getRaw(),
        .imageView = mImageView.getRaw()
    };
}

void AutoResizeImageReference::createImage()
{
    const vk::Extent2D extent = mReferenceImageRef->getExtent();
    mImageCreationArgs.extent.width = extent.width;
    mImageCreationArgs.extent.height = extent.height;
    mImageCreationArgs.extent.depth = 1;
    mImage = getOwner()->createChild<Image>(mImageCreationArgs);
    mImage->allocateMemory();
    mImageView = mImage->createImageView(mImageViewCreationArgs);
}
} // namespace iwa