#include "rendering/render_graph.hpp"

#include <sophia/util/assert.hpp>
#include <utility>

#include "vulkan/debug_label.hpp"

namespace sophia {

ImageRef RenderGraph::importImage(Image& image, Access finalAccess) {
  ImageRef ref{static_cast<u32>(this->images.size())};
  this->images.push_back({&image, finalAccess, image.getAccessState()});
  return ref;
}

void RenderGraph::addPass(std::string name, const SetupFn& setup,
                          ExecuteFn execute) {
  Pass& pass = this->passes.emplace_back();
  pass.name = std::move(name);
  pass.execute = std::move(execute);

  PassBuilder builder{*this, pass};
  setup(builder);
}

void RenderGraph::execute(vk::CommandBuffer commandBuffer) {
  for (Pass& pass : this->passes) {
    this->recordPass(commandBuffer, pass);
  }

  PendingBarriers pending;
  for (u32 i = 0; i < this->images.size(); i++) {
    this->transitionTo(ImageRef{i}, getAccessState(this->images[i].finalAccess),
                       pending);
  }
  this->flush(commandBuffer, pending);

  for (ImportedImage& imported : this->images) {
    imported.image->recordLayout(imported.state.layout);
  }

  this->passes.clear();
}

void RenderGraph::transitionTo(ImageRef ref, const AccessState& next,
                               PendingBarriers& pending) {
  SOPHIA_ASSERT(ref.isValid() && ref.index < this->images.size(),
                "ImageRef doesn't belong to this graph");

  ImportedImage& imported = this->images[ref.index];

  if (needsBarrier(imported.state, next)) {
    pending.barriers.push_back(makeImageBarrier(
        imported.image->get(), imported.image->getSubresourceRange(),
        imported.state, next));
    pending.srcStages |= imported.state.stage;
    pending.dstStages |= next.stage;
    imported.state = next;
  } else {
    // Same-layout read after read. Nothing to record, but a later write has
    // to wait on this reader too.
    imported.state.stage |= next.stage;
    imported.state.access |= next.access;
  }
}

void RenderGraph::flush(vk::CommandBuffer commandBuffer,
                        PendingBarriers& pending) {
  if (pending.barriers.empty()) {
    return;
  }

  commandBuffer.pipelineBarrier(pending.srcStages, pending.dstStages, {},
                                nullptr, nullptr, pending.barriers);
  pending = PendingBarriers{};
}

void RenderGraph::recordPass(vk::CommandBuffer commandBuffer, Pass& pass) {
  beginDebugLabel(commandBuffer, pass.name.c_str());

  PendingBarriers pending;
  for (const ImageUse& use : pass.uses) {
    this->transitionTo(use.image, getAccessState(use.access), pending);
  }
  this->flush(commandBuffer, pending);

  if (pass.colorAttachments.empty()) {
    pass.execute(commandBuffer);
  } else {
    this->recordRendering(commandBuffer, pass);
  }

  endDebugLabel(commandBuffer);
}

void RenderGraph::recordRendering(vk::CommandBuffer commandBuffer, Pass& pass) {
  std::vector<vk::RenderingAttachmentInfo> attachments;
  attachments.reserve(pass.colorAttachments.size());

  vk::Extent3D imageExtent =
      this->images[pass.colorAttachments.front().image.index]
          .image->getExtent();
  vk::Extent2D extent{imageExtent.width, imageExtent.height};

  for (const ColorAttachment& color : pass.colorAttachments) {
    const Image& image = *this->images[color.image.index].image;
    SOPHIA_ASSERT(image.getExtent().width == extent.width &&
                      image.getExtent().height == extent.height,
                  "Pass", pass.name,
                  "has color attachments of different sizes");

    vk::RenderingAttachmentInfo attachment{};
    attachment.imageView = image.getView();
    attachment.imageLayout = vk::ImageLayout::eColorAttachmentOptimal;
    attachment.loadOp = color.loadOp;
    attachment.storeOp = vk::AttachmentStoreOp::eStore;
    attachment.clearValue.color = vk::ClearColorValue{color.clearColor};
    attachments.push_back(attachment);
  }

  vk::RenderingInfo renderingInfo{};
  renderingInfo.renderArea = vk::Rect2D{{0, 0}, extent};
  renderingInfo.layerCount = 1;
  renderingInfo.colorAttachmentCount = static_cast<u32>(attachments.size());
  renderingInfo.pColorAttachments = attachments.data();

  commandBuffer.beginRendering(renderingInfo);

  vk::Viewport viewport{0.0f,
                        0.0f,
                        static_cast<float>(extent.width),
                        static_cast<float>(extent.height),
                        0.0f,
                        1.0f};
  vk::Rect2D scissor{{0, 0}, extent};
  commandBuffer.setViewport(0, viewport);
  commandBuffer.setScissor(0, scissor);

  pass.execute(commandBuffer);

  commandBuffer.endRendering();
}

RenderGraph::PassBuilder& RenderGraph::PassBuilder::writeColor(
    ImageRef image, vk::AttachmentLoadOp loadOp, ClearColor clearColor) {
  this->use(image, Access::ColorWrite);
  this->pass.colorAttachments.push_back({image, loadOp, clearColor});
  return *this;
}

RenderGraph::PassBuilder& RenderGraph::PassBuilder::sample(ImageRef image) {
  return this->use(image, Access::SampledRead);
}

RenderGraph::PassBuilder& RenderGraph::PassBuilder::use(ImageRef image,
                                                        Access access) {
  SOPHIA_ASSERT(image.isValid() && image.index < this->graph.images.size(),
                "ImageRef doesn't belong to this graph");
  for (const ImageUse& existing : this->pass.uses) {
    SOPHIA_ASSERT(existing.image.index != image.index, "Pass", this->pass.name,
                  "uses the same image twice");
  }

  this->pass.uses.push_back({image, access});
  return *this;
}

}  // namespace sophia