#include "image.hpp"

#include <cstring>
#include <sophia/util/logger.hpp>
#include <stdexcept>

#include "device.hpp"

namespace sophia {

void Image::transitionLayout(vk::CommandBuffer commandBuffer,
                             vk::ImageLayout newLayout) {
  vk::ImageMemoryBarrier barrier{};
  barrier.oldLayout = this->layout;
  barrier.newLayout = newLayout;
  barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.image = this->get();
  barrier.subresourceRange.aspectMask = this->aspectMask;
  barrier.subresourceRange.baseMipLevel = 0;
  barrier.subresourceRange.levelCount = this->mipLevels;
  barrier.subresourceRange.baseArrayLayer = 0;
  barrier.subresourceRange.layerCount = this->arrayLayers;

  vk::PipelineStageFlags srcStage;
  vk::PipelineStageFlags dstStage;

  if (this->layout == vk::ImageLayout::eUndefined &&
      newLayout == vk::ImageLayout::eTransferDstOptimal) {
    barrier.srcAccessMask = {};
    barrier.dstAccessMask = vk::AccessFlagBits::eTransferWrite;
    srcStage = vk::PipelineStageFlagBits::eTopOfPipe;
    dstStage = vk::PipelineStageFlagBits::eTransfer;
  } else if (this->layout == vk::ImageLayout::eTransferDstOptimal &&
             newLayout == vk::ImageLayout::eShaderReadOnlyOptimal) {
    barrier.srcAccessMask = vk::AccessFlagBits::eTransferWrite;
    barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead;
    srcStage = vk::PipelineStageFlagBits::eTransfer;
    dstStage = vk::PipelineStageFlagBits::eFragmentShader;
  } else if (this->layout == vk::ImageLayout::eUndefined &&
             newLayout == vk::ImageLayout::eDepthStencilAttachmentOptimal) {
    barrier.srcAccessMask = {};
    barrier.dstAccessMask = vk::AccessFlagBits::eDepthStencilAttachmentRead |
                            vk::AccessFlagBits::eDepthStencilAttachmentWrite;
    srcStage = vk::PipelineStageFlagBits::eTopOfPipe;
    dstStage = vk::PipelineStageFlagBits::eEarlyFragmentTests;
  } else if (this->layout == vk::ImageLayout::eUndefined &&
             newLayout == vk::ImageLayout::eColorAttachmentOptimal) {
    barrier.srcAccessMask = {};
    barrier.dstAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;
    srcStage = vk::PipelineStageFlagBits::eTopOfPipe;
    dstStage = vk::PipelineStageFlagBits::eColorAttachmentOutput;
  } else if (this->layout == vk::ImageLayout::eUndefined &&
             newLayout == vk::ImageLayout::eGeneral) {
    barrier.srcAccessMask = {};
    barrier.dstAccessMask =
        vk::AccessFlagBits::eShaderRead | vk::AccessFlagBits::eShaderWrite;
    srcStage = vk::PipelineStageFlagBits::eTopOfPipe;
    dstStage = vk::PipelineStageFlagBits::eComputeShader;
  } else if (this->layout == vk::ImageLayout::eColorAttachmentOptimal &&
             newLayout == vk::ImageLayout::ePresentSrcKHR) {
    barrier.srcAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;
    barrier.dstAccessMask = {};
    srcStage = vk::PipelineStageFlagBits::eColorAttachmentOutput;
    dstStage = vk::PipelineStageFlagBits::eBottomOfPipe;
  } else if (this->layout == vk::ImageLayout::eColorAttachmentOptimal &&
             newLayout == vk::ImageLayout::eShaderReadOnlyOptimal) {
    barrier.srcAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;
    barrier.dstAccessMask = vk::AccessFlagBits::eShaderRead;
    srcStage = vk::PipelineStageFlagBits::eColorAttachmentOutput;
    dstStage = vk::PipelineStageFlagBits::eFragmentShader;
  } else if (this->layout == vk::ImageLayout::eShaderReadOnlyOptimal &&
             newLayout == vk::ImageLayout::eColorAttachmentOptimal) {
    barrier.srcAccessMask = vk::AccessFlagBits::eShaderRead;
    barrier.dstAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;
    srcStage = vk::PipelineStageFlagBits::eFragmentShader;
    dstStage = vk::PipelineStageFlagBits::eColorAttachmentOutput;
  } else if (this->layout == vk::ImageLayout::ePresentSrcKHR &&
             newLayout == vk::ImageLayout::eColorAttachmentOptimal) {
    barrier.srcAccessMask = {};
    barrier.dstAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;
    srcStage = vk::PipelineStageFlagBits::eTopOfPipe;
    dstStage = vk::PipelineStageFlagBits::eColorAttachmentOutput;
  } else {
    log::fatal("Unsupported image layout transition.");
    throw std::runtime_error("Unsupported image layout transition.");
  }

  commandBuffer.pipelineBarrier(srcStage, dstStage, {}, nullptr, nullptr,
                                barrier);

  this->layout = newLayout;
}

void Image::transitionLayoutImmediate(vk::ImageLayout newLayout) {
  vk::CommandBuffer commandBuffer = this->device.beginSingleTimeCommands();
  this->transitionLayout(commandBuffer, newLayout);
  this->device.endSingleTimeCommands(commandBuffer);
}

Image::Builder& Image::Builder::setExtent(u32 width, u32 height, u32 depth) {
  this->extent = vk::Extent3D{width, height, depth};
  return *this;
}

Image::Builder& Image::Builder::setFormat(vk::Format format) {
  this->format = format;
  return *this;
}

Image::Builder& Image::Builder::setUsage(vk::ImageUsageFlags usage) {
  this->usage = usage;
  return *this;
}

Image::Builder& Image::Builder::setTiling(vk::ImageTiling tiling) {
  this->tiling = tiling;
  return *this;
}

Image::Builder& Image::Builder::setMemoryProperties(
    vk::MemoryPropertyFlags properties) {
  this->memoryProperties = properties;
  return *this;
}

Image::Builder& Image::Builder::setSamples(vk::SampleCountFlagBits samples) {
  this->samples = samples;
  return *this;
}

Image::Builder& Image::Builder::setImageType(vk::ImageType imageType) {
  this->imageType = imageType;
  return *this;
}

Image::Builder& Image::Builder::setMipLevels(u32 mipLevels) {
  this->mipLevels = mipLevels;
  return *this;
}

Image::Builder& Image::Builder::setArrayLayers(u32 arrayLayers) {
  this->arrayLayers = arrayLayers;
  return *this;
}

Image::Builder& Image::Builder::setAspectMask(vk::ImageAspectFlags aspectMask) {
  this->aspectMask = aspectMask;
  return *this;
}

Image::Builder& Image::Builder::asComputeStorage() {
  // Storage images commonly can't use sRGB formats without an extension, so
  // default to a plain UNORM format; override with setFormat() if needed.
  this->format = vk::Format::eR8G8B8A8Unorm;
  this->usage = vk::ImageUsageFlagBits::eStorage;
  this->aspectMask = vk::ImageAspectFlagBits::eColor;
  this->tiling = vk::ImageTiling::eOptimal;
  this->memoryProperties = vk::MemoryPropertyFlagBits::eDeviceLocal;
  return *this;
}

Image::Builder& Image::Builder::asColorAttachment() {
  this->usage = vk::ImageUsageFlagBits::eColorAttachment |
                vk::ImageUsageFlagBits::eSampled;
  this->aspectMask = vk::ImageAspectFlagBits::eColor;
  this->tiling = vk::ImageTiling::eOptimal;
  this->memoryProperties = vk::MemoryPropertyFlagBits::eDeviceLocal;
  return *this;
}

Image::Builder& Image::Builder::asDepthAttachment() {
  this->format = vk::Format::eD32Sfloat;
  this->usage = vk::ImageUsageFlagBits::eDepthStencilAttachment;
  this->aspectMask = vk::ImageAspectFlagBits::eDepth;
  this->tiling = vk::ImageTiling::eOptimal;
  this->memoryProperties = vk::MemoryPropertyFlagBits::eDeviceLocal;
  return *this;
}

Image::Builder& Image::Builder::wrapExisting(vk::Image image) {
  this->isExternalImage = true;
  this->externalImage = image;
  return *this;
}

Image::Builder& Image::Builder::setInitialData(const void* data,
                                               vk::DeviceSize size) {
  this->initialData = data;
  this->initialDataSize = size;
  return *this;
}

Image::Builder& Image::Builder::withView(bool enabled) {
  this->createView = enabled;
  return *this;
}

Image::Builder& Image::Builder::withSampler(
    vk::Filter filter, vk::SamplerAddressMode addressMode) {
  this->createSampler = true;
  this->samplerFilter = filter;
  this->samplerAddressMode = addressMode;
  return *this;
}

std::unique_ptr<Image> Image::Builder::build() {
  vk::ImageUsageFlags usage = this->usage;

  auto result = std::unique_ptr<Image>(new Image(this->device));
  result->format = this->format;
  result->extent = this->extent;
  result->aspectMask = this->aspectMask;
  result->mipLevels = this->mipLevels;
  result->arrayLayers = this->arrayLayers;

  if (this->isExternalImage) {
    if (this->initialData != nullptr) {
      log::warning(
          "Image::Builder::setInitialData is ignored when wrapExisting is "
          "used.");
    }
    result->externalImage = this->externalImage;
    // Wrapped images (e.g. swapchain images) are handed to us already in a
    // known layout controlled by their owner; we can't observe it here, so
    // leave `layout` at eUndefined and let the caller call
    // transitionLayout() to bring it in sync before relying on it.
  } else {
    if (this->initialData != nullptr) {
      usage |= vk::ImageUsageFlagBits::eTransferDst;
    }
    if (this->createView || this->createSampler) {
      usage |= vk::ImageUsageFlagBits::eSampled;
    }

    vk::ImageCreateInfo imageInfo{};
    imageInfo.imageType = this->imageType;
    imageInfo.extent = this->extent;
    imageInfo.mipLevels = this->mipLevels;
    imageInfo.arrayLayers = this->arrayLayers;
    imageInfo.format = this->format;
    imageInfo.tiling = this->tiling;
    imageInfo.initialLayout = vk::ImageLayout::eUndefined;
    imageInfo.usage = usage;
    imageInfo.sharingMode = vk::SharingMode::eExclusive;
    imageInfo.samples = this->samples;

    result->image = this->device.get()->createImageUnique(imageInfo);

    vk::MemoryRequirements memoryRequirements =
        this->device.get()->getImageMemoryRequirements(result->image.get());

    vk::MemoryAllocateInfo allocInfo{};
    allocInfo.allocationSize = memoryRequirements.size;
    allocInfo.memoryTypeIndex = this->device.findMemoryType(
        memoryRequirements.memoryTypeBits, this->memoryProperties);

    result->memory = this->device.get()->allocateMemoryUnique(allocInfo);
    this->device.get()->bindImageMemory(result->image.get(),
                                        result->memory.get(), 0);
  }

  if (!this->isExternalImage && this->initialData != nullptr) {
    vk::Buffer stagingBuffer;
    vk::DeviceMemory stagingMemory;
    this->device.createBuffer(this->initialDataSize,
                              vk::BufferUsageFlagBits::eTransferSrc,
                              vk::MemoryPropertyFlagBits::eHostVisible |
                                  vk::MemoryPropertyFlagBits::eHostCoherent,
                              stagingBuffer, stagingMemory);

    void* mapped =
        this->device.get()->mapMemory(stagingMemory, 0, this->initialDataSize);
    memcpy(mapped, this->initialData,
           static_cast<size_t>(this->initialDataSize));
    this->device.get()->unmapMemory(stagingMemory);

    vk::BufferImageCopy region{};
    region.bufferOffset = 0;
    region.bufferRowLength = 0;
    region.bufferImageHeight = 0;
    region.imageSubresource.aspectMask = this->aspectMask;
    region.imageSubresource.mipLevel = 0;
    region.imageSubresource.baseArrayLayer = 0;
    region.imageSubresource.layerCount = this->arrayLayers;
    region.imageOffset = vk::Offset3D{0, 0, 0};
    region.imageExtent = this->extent;

    vk::CommandBuffer commandBuffer = this->device.beginSingleTimeCommands();
    result->transitionLayout(commandBuffer,
                             vk::ImageLayout::eTransferDstOptimal);
    commandBuffer.copyBufferToImage(stagingBuffer, result->get(),
                                    vk::ImageLayout::eTransferDstOptimal,
                                    region);
    if (usage & vk::ImageUsageFlagBits::eSampled) {
      result->transitionLayout(commandBuffer,
                               vk::ImageLayout::eShaderReadOnlyOptimal);
    }
    this->device.endSingleTimeCommands(commandBuffer);

    this->device.get()->destroyBuffer(stagingBuffer);
    this->device.get()->freeMemory(stagingMemory);
  }

  if (this->createView) {
    vk::ImageViewCreateInfo viewInfo{};
    viewInfo.image = result->get();
    viewInfo.viewType = vk::ImageViewType::e2D;
    viewInfo.format = this->format;
    viewInfo.subresourceRange.aspectMask = this->aspectMask;
    viewInfo.subresourceRange.baseMipLevel = 0;
    viewInfo.subresourceRange.levelCount = this->mipLevels;
    viewInfo.subresourceRange.baseArrayLayer = 0;
    viewInfo.subresourceRange.layerCount = this->arrayLayers;

    result->view = this->device.get()->createImageViewUnique(viewInfo);
  }

  if (this->createSampler) {
    vk::SamplerCreateInfo samplerInfo{};
    samplerInfo.magFilter = this->samplerFilter;
    samplerInfo.minFilter = this->samplerFilter;
    samplerInfo.addressModeU = this->samplerAddressMode;
    samplerInfo.addressModeV = this->samplerAddressMode;
    samplerInfo.addressModeW = this->samplerAddressMode;
    samplerInfo.anisotropyEnable = VK_FALSE;
    samplerInfo.maxAnisotropy = 1.0f;
    samplerInfo.borderColor = vk::BorderColor::eIntOpaqueBlack;
    samplerInfo.unnormalizedCoordinates = VK_FALSE;
    samplerInfo.compareEnable = VK_FALSE;
    samplerInfo.compareOp = vk::CompareOp::eAlways;
    samplerInfo.mipmapMode = vk::SamplerMipmapMode::eLinear;
    samplerInfo.minLod = 0.0f;
    samplerInfo.maxLod = static_cast<float>(this->mipLevels);

    result->sampler = this->device.get()->createSamplerUnique(samplerInfo);
  }

  return result;
}

}  // namespace sophia