#include "Images.hpp"
#include "Buffers.hpp"
#include "../Allocator.hpp"
#include "../../../../core/utils/logger.hpp"
#include <algorithm>

using namespace render::vk;

Image::~Image() {
    vkDestroyImage(Allocator::GetDefault()->getDevice(), ref, ALLOC);
    vkDestroyImageView(Allocator::GetDefault()->getDevice(), view, ALLOC);
}
Texture::~Texture() {
    vkDestroySampler(Allocator::GetDefault()->getDevice(), sampler, ALLOC);
}

memory::ptr createImage(const Image::requirement& req, VkMemoryPropertyFlags properties, const std::vector<render::data_view>& datas, Image::info& out) {
    auto alloc = Allocator::GetDefault();
    auto device = alloc->getDevice();

    VkImageCreateInfo info{};
    info.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    info.imageType = VK_IMAGE_TYPE_2D;
    info.extent.width = req.size.width;
    info.extent.height = req.size.height;
    info.extent.depth = 1;
    info.mipLevels = req.mipmapLevels;
    info.arrayLayers = req.layers;
    info.format = static_cast<VkFormat>(req.format);
    info.tiling = req.optimal ? VK_IMAGE_TILING_OPTIMAL : VK_IMAGE_TILING_LINEAR;
    info.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    info.usage = static_cast<VkImageUsageFlags>(req.usage);
    info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    info.samples = static_cast<VkSampleCountFlagBits>(req.samples);
    info.flags = 0;
    if (req.cube) {
        info.flags |= VK_IMAGE_CREATE_CUBE_COMPATIBLE_BIT;
    }

    if (!datas.empty()) {
        info.usage |= VK_IMAGE_USAGE_TRANSFER_DST_BIT;
    }

    if (vkCreateImage(device, &info, ALLOC, &out.ref) != VK_SUCCESS) {
        LOG_E("Failed to create image");
        return memory::GetNull();
    }
    out.offset = 0;

    VkMemoryRequirements memRequirements;
    vkGetImageMemoryRequirements(device, out.ref, &memRequirements);

    auto memory = alloc->allocate(memRequirements, properties, true);
    if (!memory || vkBindImageMemory(device, out.ref, memory->ref, memory->offset) != VK_SUCCESS) {
        LOG_E("Failed to allocate image memory");
        return memory::GetNull();
    }

    if (!datas.empty()) {
        assert(datas.size() <= req.layers);
        const auto maxSize = [&] {
            size_t max = 0;
            for (auto& data: datas)
                max = std::max(max, data.size);
            return max;
        }();
        if(auto staging = WritableBuffer::Create(maxSize)) {
            alloc->transitionImageLayout(out.ref, info.format, info.initialLayout, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, req.mipmapLevels, req.layers);
            info.initialLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
            for (size_t layer = 0; layer < req.layers && layer < datas.size(); layer++) {
                if(datas[layer].isUsable()) {
                    staging->write(datas[layer], 0);
                    alloc->copyBufferToImage(staging->getRef(), out.ref, info.extent.width, info.extent.height, req.mipmapLevels, layer);
                }
            }
        } else {
            LOG_E("Cannot allocate staging memory");
            return memory::GetNull();
        }
    }
    alloc->transitionImageLayout(out.ref, info.format, info.initialLayout, static_cast<VkImageLayout>(req.layout), req.mipmapLevels, req.layers);

    VkImageViewCreateInfo viewInfo{};
    viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    viewInfo.image = out.ref;
    if (req.cube) {
        viewInfo.viewType = VK_IMAGE_VIEW_TYPE_CUBE;
    } else if (req.layers > 1) {
        viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D_ARRAY;
    } else {
        viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
    }
    viewInfo.format = static_cast<VkFormat>(req.format);
    viewInfo.subresourceRange.aspectMask = static_cast<VkImageAspectFlags>(req.aspect);
    viewInfo.subresourceRange.baseMipLevel = 0;
    viewInfo.subresourceRange.levelCount = req.mipmapLevels;
    viewInfo.subresourceRange.baseArrayLayer = 0;
    viewInfo.subresourceRange.layerCount = req.layers;

    if (vkCreateImageView(device, &viewInfo, ALLOC, &out.view) != VK_SUCCESS) {
        LOG_E("Failed to create texture image view!");
        return memory::GetNull();
    }
    return memory;
}

std::unique_ptr<Image> Image::Create(const requirement & req) {
    vk::Image::info img;
    auto mem = createImage(req, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, {}, img);
    if(!mem) {
        FATAL("Cannot create texture image");
    }
    return std::unique_ptr<Image>(new Image(img.ref, img.view, std::move(mem)));
}

VkSampler createSampler(const Texture::sampling& props, uint32_t mipmapLevels) {
    VkSamplerCreateInfo samplerInfo{};
    samplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
    samplerInfo.magFilter = props.magLinear ? VK_FILTER_LINEAR : VK_FILTER_NEAREST;
    samplerInfo.minFilter = props.minLinear ? VK_FILTER_LINEAR : VK_FILTER_NEAREST;

    const auto wrap = static_cast<VkSamplerAddressMode>(props.wrap);
    samplerInfo.addressModeU = wrap;
    samplerInfo.addressModeV = wrap;
    samplerInfo.addressModeW = wrap;

    auto maxAnisotropy = Allocator::GetDefault()->getCapabilities().maxAnisotropy;
    if (maxAnisotropy && props.anisotropy > 0) {
        samplerInfo.anisotropyEnable = VK_TRUE;
        samplerInfo.maxAnisotropy = std::min<float>(maxAnisotropy.value(), 1 << (props.anisotropy-1));
    } else {
        samplerInfo.anisotropyEnable = VK_FALSE;
        samplerInfo.maxAnisotropy = 1.f;
    }

    samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
    samplerInfo.unnormalizedCoordinates = VK_FALSE;

    samplerInfo.compareEnable = VK_FALSE;
    samplerInfo.compareOp = VK_COMPARE_OP_ALWAYS;

    samplerInfo.mipmapMode = props.minLinear ? VK_SAMPLER_MIPMAP_MODE_LINEAR : VK_SAMPLER_MIPMAP_MODE_NEAREST;
    samplerInfo.mipLodBias = std::min(Allocator::GetDefault()->getCapabilities().maxLodBias, props.mipmapLod);
    samplerInfo.minLod = 0.0f;
    samplerInfo.maxLod = props.mipmap ? mipmapLevels : 0.0f;

    if (props.mipmap && mipmapLevels <= 1) {
        LOG_D("Sampler requires mipmap but image does not");
    }

    VkSampler sampler;
    if (vkCreateSampler(Allocator::GetDefault()->getDevice(), &samplerInfo, ALLOC, &sampler) != VK_SUCCESS) {
        FATAL("Failed to create texture sampler!");
    }
    return sampler;
}

std::unique_ptr<Texture> Texture::LoadFromFile(const std::string& path, const sampling& props) {
    std::vector<unsigned char> data;
    auto header = [&] {
        if (auto header = render::Image::Read(path, data)) {
            return header.value();
        }
        FATAL("Cannot read texture");
    }();

    vk::Image::info img;
    auto mem = createImage(requirement::Texture(header),
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, {data}, img);
    if(!mem) {
        FATAL("Cannot create texture image");
    }

    return std::unique_ptr<Texture>(new Texture(createSampler(props, header.mipmapLevels), img.view, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, img.ref, std::move(mem)));
}

std::unique_ptr<TextureCube> TextureCube::LoadFromFiles(const std::array<std::string, 6>& paths, const sampling& props) {
    std::vector<std::vector<unsigned char>> datas;
    std::vector<data_view> views;
    views.reserve(6);
    datas.resize(1);
    auto header = [&] {
        if (auto header = render::Image::Read(paths.at(0), datas.at(0))) {
            views.push_back(datas.at(0));
            return header.value();
        }
        FATAL("Cannot read first texture");
    }();
    datas.resize(paths.size());
    for (size_t i = 1; i < paths.size(); i++) {
        if(!render::Image::Read(paths.at(i), datas.at(i)).has_value()) {
            FATAL("Cannot read depth texture");
        }
        views.push_back(datas.at(i));
    }

    vk::Image::info img;
    auto mem = createImage(requirement::Texture(header, paths.size(), true),
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, views, img);
    if(!mem) {
        FATAL("Cannot create texture cube image");
    }

    return std::unique_ptr<TextureCube>(new TextureCube(createSampler(props, header.mipmapLevels), img.view, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, img.ref, std::move(mem)));
}
std::unique_ptr<TextureCube> TextureCube::LoadFromFiles(const std::string& prefix, const sampling& props) {
    const std::array<std::string, 6> faces {
        "right",
        "left",
        "top",
        "bottom",
        "front",
        "back"
    };
    std::array<std::string, 6> paths;
    std::transform(faces.begin(), faces.end(), paths.begin(),
        [prefix](const std::string &face) -> std::string { return prefix + "." + face + ".dds"; });
    return LoadFromFiles(paths, props);
}

std::unique_ptr<TextureArray> TextureArray::LoadFromFiles(const std::vector<std::string>& paths, const sampling& props, bool srgb) {
    std::vector<std::vector<unsigned char>> datas;
    std::vector<data_view> views;
    datas.resize(paths.size());
    views.reserve(paths.size());
    auto header = [&] {
        if (auto header = render::Image::Read(paths.at(0), datas.at(0), srgb)) {
            views.push_back(datas.at(0));
            return header.value();
        }
        FATAL("Cannot read first texture");
    }();
    for (size_t i = 1; i < paths.size(); i++) {
        if(!render::Image::Read(paths.at(i), datas.at(i), srgb).has_value()) {
            FATAL("Cannot read depth texture");
        }
        views.push_back(datas.at(i));
    }

    vk::Image::info img;
    auto mem = createImage(requirement::Texture(header, paths.size()),
        VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, views, img);
    if(!mem) {
        FATAL("Cannot create texture cube image");
    }

    return std::unique_ptr<TextureArray>(new TextureArray(paths.size(), createSampler(props, header.mipmapLevels), img.view, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, img.ref, std::move(mem)));
}