#include "texture.h"


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


texture::texture() {
    memory = nullptr;
}

void texture::upload(const void* data, size_t size) {
    renderer* r = application::get()->get_renderer();
    const vk::Device& device = r->device;

    // Create Upload Buffer
    vk::BufferCreateInfo buffer_info;
    buffer_info.setSize(size);
    buffer_info.setUsage(vk::BufferUsageFlagBits::eTransferSrc);
    buffer_info.setSharingMode(vk::SharingMode::eExclusive);
    vk::Buffer upload_buffer = device.createBuffer(buffer_info);

    vk::MemoryRequirements req = device.getBufferMemoryRequirements(upload_buffer);
    uint32_t upload_memory_type = vk::su::findMemoryType(r->physical_device.getMemoryProperties(), req.memoryTypeBits,
                                                         vk::MemoryPropertyFlagBits::eHostVisible);
    vk::MemoryAllocateInfo alloc_info;
    alloc_info.setAllocationSize(req.size);
    alloc_info.setMemoryTypeIndex(upload_memory_type);
    vk::DeviceMemory upload_memory = device.allocateMemory(alloc_info);
    device.bindBufferMemory(upload_buffer, upload_memory, 0);

    void* p = device.mapMemory(upload_memory, 0, size);
    memcpy(p, data, size);
    vk::MappedMemoryRange memory_range;
    memory_range.setMemory(upload_memory);
    memory_range.setSize(size);
    device.flushMappedMemoryRanges(memory_range);
    device.unmapMemory(upload_memory);

    const vk::CommandBuffer command_buffer = r->create_command_buffer(vk::CommandBufferLevel::ePrimary, true);

    vk::ImageMemoryBarrier copy_barrier;
    copy_barrier.setDstAccessMask(vk::AccessFlagBits::eTransferWrite);
    copy_barrier.setOldLayout(vk::ImageLayout::eUndefined);
    copy_barrier.setNewLayout(vk::ImageLayout::eTransferDstOptimal);
    copy_barrier.setSrcQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED);
    copy_barrier.setDstQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED);
    copy_barrier.setImage(image);
    copy_barrier.setSubresourceRange({vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1});
    command_buffer.pipelineBarrier(vk::PipelineStageFlagBits::eHost, vk::PipelineStageFlagBits::eTransfer, {}, {}, {},
                                   copy_barrier);

    vk::BufferImageCopy region;
    region.setImageSubresource({vk::ImageAspectFlagBits::eColor, 0, 0, 1});
    region.setImageExtent({static_cast<uint32_t>(width_), static_cast<uint32_t>(height_), 1});
    command_buffer.copyBufferToImage(upload_buffer, image, vk::ImageLayout::eTransferDstOptimal, region);

    vk::ImageMemoryBarrier use_barrier;
    use_barrier.setSrcAccessMask(vk::AccessFlagBits::eTransferWrite);
    use_barrier.setDstAccessMask(vk::AccessFlagBits::eShaderRead);
    use_barrier.setOldLayout(vk::ImageLayout::eTransferDstOptimal);
    use_barrier.setNewLayout(vk::ImageLayout::eShaderReadOnlyOptimal);
    use_barrier.setSrcQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED);
    use_barrier.setDstQueueFamilyIndex(VK_QUEUE_FAMILY_IGNORED);
    use_barrier.setImage(image);
    use_barrier.setSubresourceRange({vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1});
    command_buffer.pipelineBarrier(vk::PipelineStageFlagBits::eTransfer, vk::PipelineStageFlagBits::eFragmentShader, {},
                                   {}, {}, use_barrier);

    r->end_command_buffer(command_buffer);
    device.waitIdle();

    device.destroyBuffer(upload_buffer);
    device.freeMemory(upload_memory);

    // device.freeCommandBuffers(command_pool, command_buffer);
}

void texture::on_init() {
    const renderer* r = application::get()->get_renderer();
    const vk::Device& device = r->device;
    const vk::PhysicalDevice& physical_device = r->physical_device;

    vk::ImageCreateInfo image_info;
    image_info.setImageType(vk::ImageType::e2D);
    image_info.setFormat(format);
    image_info.setExtent({ width_, height_, 1 });
    image_info.setMipLevels(1);
    image_info.setArrayLayers(1);
    image_info.setSamples(vk::SampleCountFlagBits::e1);
    image_info.setTiling(vk::ImageTiling::eOptimal);
    image_info.setUsage(vk::ImageUsageFlagBits::eSampled | vk::ImageUsageFlagBits::eTransferDst | vk::ImageUsageFlagBits::eStorage);
    image_info.setSharingMode(vk::SharingMode::eExclusive);
    image_info.setInitialLayout(vk::ImageLayout::eUndefined);
    image = device.createImage(image_info);

    const vk::PhysicalDeviceMemoryProperties memory_properties = physical_device.getMemoryProperties();
    const 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 view_info;
    view_info.setImage(image);
    view_info.setViewType(vk::ImageViewType::e2D);
    view_info.setFormat(format);
    view_info.setSubresourceRange({vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1});
    image_view = device.createImageView(view_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);
}