#include "render_target.h"

#include "renderer.h"
#include "application/application.h"
#include "utils/utils.hpp"

void* render_target::lock(lock_state state) const {
    const auto r = application::get()->get_renderer();
    const auto& device = r->device;

    vk::ImageSubresource subresource;
    subresource.setAspectMask(vk::ImageAspectFlagBits::eColor);
    subresource.setMipLevel(0);
    subresource.setArrayLayer(0);

    const vk::SubresourceLayout subresource_layout = device.getImageSubresourceLayout(image, subresource);

    void* data;
    const vk::Result result = device.mapMemory(memory, 0, VK_WHOLE_SIZE, vk::MemoryMapFlags(), &data);
    check_vk_result(result);

    return static_cast<uint8_t*>(data) + subresource_layout.offset;
}

void render_target::unlock() {
    const auto r = application::get()->get_renderer();
    const auto& device = r->device;

    device.unmapMemory(memory);
}

void render_target::destroy() {
    render_resource::destroy();
    const auto r = application::get()->get_renderer();
    const auto& device = r->device;
    device.destroyFramebuffer(framebuffer);
}

void render_target::create(uint32_t width, uint32_t height) {
    auto r = application::get()->get_renderer();
    const auto& device = r->device;
    const auto& physical_device = r->physical_device;

    vk::ImageCreateInfo image_create_info;
    image_create_info.setImageType(vk::ImageType::e2D);
    image_create_info.setFormat(format);
    image_create_info.setExtent({static_cast<uint32_t>(width), static_cast<uint32_t>(height), 1});
    image_create_info.setMipLevels(1);
    image_create_info.setArrayLayers(1);
    image_create_info.setSamples(vk::SampleCountFlagBits::e1);
    image_create_info.setTiling(vk::ImageTiling::eOptimal);
    image_create_info.setUsage(vk::ImageUsageFlagBits::eSampled | vk::ImageUsageFlagBits::eTransferDst | vk::ImageUsageFlagBits::eColorAttachment | vk::ImageUsageFlagBits::eStorage);
    image_create_info.setSharingMode(vk::SharingMode::eExclusive);
    image_create_info.setInitialLayout(vk::ImageLayout::eUndefined);
    image = device.createImage(image_create_info);

    vk::PhysicalDeviceMemoryProperties memory_properties = physical_device.getMemoryProperties();
    vk::MemoryRequirements memory_requirements = device.getImageMemoryRequirements(image);
    const uint32_t memory_type = vk::su::findMemoryType(memory_properties, memory_requirements.memoryTypeBits, vk::MemoryPropertyFlagBits::eDeviceLocal);

    vk::MemoryAllocateInfo memory_allocate_info;
    memory_allocate_info.setAllocationSize(memory_requirements.size);
    memory_allocate_info.setMemoryTypeIndex(memory_type);
    memory = device.allocateMemory(memory_allocate_info);
    device.bindImageMemory(image, memory, 0);

    vk::ImageViewCreateInfo image_view_create_info;
    image_view_create_info.setImage(image);
    image_view_create_info.setViewType(vk::ImageViewType::e2D);
    image_view_create_info.setFormat(format);
    image_view_create_info.setSubresourceRange({vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1});
    image_view = device.createImageView(image_view_create_info);

    vk::FramebufferCreateInfo framebuffer_create_info;
    framebuffer_create_info.setRenderPass(r->main_window_data.RenderPass);
    framebuffer_create_info.setAttachmentCount(1);
    framebuffer_create_info.setAttachments(image_view);
    framebuffer_create_info.setWidth(width);
    framebuffer_create_info.setHeight(height);
    framebuffer_create_info.setLayers(1);
    framebuffer = device.createFramebuffer(framebuffer_create_info);

    vk::SamplerCreateInfo sampler_info;
    sampler_info.setMagFilter(vk::Filter::eLinear);
    sampler_info.setMinFilter(vk::Filter::eLinear);
    sampler_info.setMipmapMode(vk::SamplerMipmapMode::eLinear);
    sampler_info.setAddressModeU(vk::SamplerAddressMode::eRepeat);
    sampler_info.setAddressModeV(vk::SamplerAddressMode::eRepeat);
    sampler_info.setAddressModeW(vk::SamplerAddressMode::eRepeat);
    sampler_info.setMinLod(-1000);
    sampler_info.setMaxLod(1000);
    sampler_info.setMaxAnisotropy(1);
    sampler = device.createSampler(sampler_info);

    descriptor_set = ImGui_ImplVulkan_AddTexture(sampler, image_view, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
}

void render_target::on_init() {
    create(width_, height_);
}

void render_target::on_resize(uint32_t width, uint32_t height) {
    destroy();
    create(width, height);
}