#pragma once

#include <memory>
#include <sophia/types.hpp>
#include <vulkan/vulkan.hpp>

#include "vulkan/access.hpp"

namespace sophia {

class Device;

class Image {
 public:
  class Builder;

  Image(const Image&) = delete;
  Image& operator=(const Image&) = delete;
  Image(Image&&) = default;
  Image& operator=(Image&&) = delete;

  vk::Image get() const {
    return this->image ? this->image.get() : this->externalImage;
  }
  vk::ImageView getView() const { return this->view.get(); }
  vk::Sampler getSampler() const { return this->sampler.get(); }

  bool isExternal() const { return !this->image; }

  vk::Format getFormat() const { return this->format; }
  vk::Extent3D getExtent() const { return this->extent; }
  vk::ImageLayout getLayout() const { return this->layout; }
  vk::ImageSubresourceRange getSubresourceRange() const;

  /**
   * The layout this image was last left in, plus the stages and accesses that
   * a barrier out of it has to wait on.
   */
  AccessState getAccessState() const;

  /**
   * Updates the tracked layout without recording a barrier. For code that
   * does its own transitions (the RenderGraph) and needs to hand the final
   * layout back.
   */
  void recordLayout(vk::ImageLayout newLayout) { this->layout = newLayout; }

  void transitionLayout(vk::CommandBuffer commandBuffer,
                        vk::ImageLayout newLayout);

  void transitionLayoutImmediate(vk::ImageLayout newLayout);

 private:
  Image(Device& device) : device{device} {}

  Device& device;

  vk::Image externalImage;
  vk::UniqueImage image;
  vk::UniqueDeviceMemory memory;
  vk::UniqueImageView view;
  vk::UniqueSampler sampler;

  vk::Format format = vk::Format::eUndefined;
  vk::Extent3D extent{};
  vk::ImageAspectFlags aspectMask = vk::ImageAspectFlagBits::eColor;
  u32 mipLevels = 1;
  u32 arrayLayers = 1;
  vk::ImageLayout layout = vk::ImageLayout::eUndefined;
};

class Image::Builder {
 public:
  Builder(Device& device) : device{device} {}

  Builder& setExtent(u32 width, u32 height, u32 depth = 1);
  Builder& setFormat(vk::Format format);
  Builder& setUsage(vk::ImageUsageFlags usage);
  Builder& setTiling(vk::ImageTiling tiling);
  Builder& setMemoryProperties(vk::MemoryPropertyFlags properties);
  Builder& setSamples(vk::SampleCountFlagBits samples);
  Builder& setImageType(vk::ImageType imageType);
  Builder& setMipLevels(u32 mipLevels);
  Builder& setArrayLayers(u32 arrayLayers);
  Builder& setAspectMask(vk::ImageAspectFlags aspectMask);

  // High-level usage presets for GPGPU + Graphics. Each sets usage/aspect/
  // tiling/memory-property defaults for the common case; call the
  // individual setters afterward to override anything preset-specific.
  Builder& asComputeStorage();
  Builder& asColorAttachment();
  Builder& asDepthAttachment();

  // Wraps an image this Builder does not own instead of creating one —
  // e.g. a swapchain image. build() will skip image/memory creation and
  // teardown entirely; setUsage/setTiling/setMemoryProperties/setSamples/
  // setInitialData are ignored, but setFormat/setExtent/setAspectMask still
  // matter, since view/sampler creation relies on them.
  Builder& wrapExisting(vk::Image image);

  // Uploads `size` bytes to the image via a staging buffer once it is
  // created, transitioning it to eTransferDstOptimal (and on to
  // eShaderReadOnlyOptimal if the image is sampled) in the process.
  Builder& setInitialData(const void* data, vk::DeviceSize size);

  Builder& withView(bool enabled = true);
  Builder& withSampler(
      vk::Filter filter = vk::Filter::eLinear,
      vk::SamplerAddressMode addressMode = vk::SamplerAddressMode::eRepeat);

  std::unique_ptr<Image> build();

 private:
  Device& device;

  vk::Extent3D extent{1, 1, 1};
  vk::Format format = vk::Format::eR8G8B8A8Srgb;
  vk::ImageUsageFlags usage = vk::ImageUsageFlagBits::eSampled;
  vk::ImageTiling tiling = vk::ImageTiling::eOptimal;
  vk::MemoryPropertyFlags memoryProperties =
      vk::MemoryPropertyFlagBits::eDeviceLocal;
  vk::SampleCountFlagBits samples = vk::SampleCountFlagBits::e1;
  vk::ImageType imageType = vk::ImageType::e2D;
  vk::ImageAspectFlags aspectMask = vk::ImageAspectFlagBits::eColor;
  u32 mipLevels = 1;
  u32 arrayLayers = 1;

  const void* initialData = nullptr;
  vk::DeviceSize initialDataSize = 0;

  bool isExternalImage = false;
  vk::Image externalImage;

  bool createView = false;
  bool createSampler = false;
  vk::Filter samplerFilter = vk::Filter::eLinear;
  vk::SamplerAddressMode samplerAddressMode = vk::SamplerAddressMode::eRepeat;
};

}  // namespace sophia