#pragma once

#include <array>
#include <functional>
#include <sophia/types.hpp>
#include <string>
#include <vector>
#include <vulkan/vulkan.hpp>

#include "vulkan/access.hpp"
#include "vulkan/image.hpp"

namespace sophia {

/**
 * Names an image within one RenderGraph. Only valid for the graph that
 * returned it.
 */
struct ImageRef {
  u32 index = INVALID_INDEX;

  static constexpr u32 INVALID_INDEX = 0xFFFFFFFF;

  bool isValid() const { return this->index != INVALID_INDEX; }
};

using ClearColor = std::array<float, 4>;

/**
 * A frame's passes and the images they touch, recorded into a single command
 * buffer. Passes run in the order they were added; the graph's job is to
 * work out the barriers and layout transitions between them from what each
 * pass says it does to each image, and to begin/end rendering for passes that
 * write color attachments.
 *
 * Built fresh every frame:
 *
 *   RenderGraph graph;
 *   ImageRef backbuffer = graph.importImage(swapchainImage, Access::Present);
 *   graph.addPass(
 *       "main", [&](PassBuilder& pass) { pass.writeColor(backbuffer); },
 *       [&](vk::CommandBuffer cb) { draw(cb); });
 *   graph.execute(commandBuffer);
 *
 * Imported images start in whatever layout Image reports and are left in
 * their final access when execute() returns, with the Image's tracked layout
 * updated to match.
 */
class RenderGraph {
 public:
  class PassBuilder;

  using SetupFn = std::function<void(PassBuilder&)>;
  using ExecuteFn = std::function<void(vk::CommandBuffer)>;

  RenderGraph() = default;

  RenderGraph(const RenderGraph&) = delete;
  RenderGraph& operator=(const RenderGraph&) = delete;

  /**
   * Brings an image the graph doesn't own into it. It must outlive
   * execute(). After the last pass it is transitioned to `finalAccess`.
   */
  ImageRef importImage(Image& image, Access finalAccess);

  /**
   * Adds a pass. `setup` runs immediately and declares what the pass does to
   * images; `execute` runs later, from execute(), with the barriers already
   * recorded and (if the pass writes color) rendering begun. `name` labels
   * the pass's commands in debugging tools such as RenderDoc.
   */
  void addPass(std::string name, const SetupFn& setup, ExecuteFn execute);

  /**
   * Records every pass, and the barriers between them, into `commandBuffer`.
   */
  void execute(vk::CommandBuffer commandBuffer);

 private:
  struct ImageUse {
    ImageRef image;
    Access access;
  };

  struct ColorAttachment {
    ImageRef image;
    vk::AttachmentLoadOp loadOp;
    ClearColor clearColor;
  };

  struct Pass {
    std::string name;
    std::vector<ImageUse> uses;
    std::vector<ColorAttachment> colorAttachments;
    ExecuteFn execute;
  };

  struct ImportedImage {
    Image* image;
    Access finalAccess;
    // What the image was last left in, as of the barriers recorded so far.
    AccessState state;
  };

  struct PendingBarriers {
    std::vector<vk::ImageMemoryBarrier> barriers;
    vk::PipelineStageFlags srcStages;
    vk::PipelineStageFlags dstStages;
  };

  void transitionTo(ImageRef ref, const AccessState& next,
                    PendingBarriers& pending);
  void flush(vk::CommandBuffer commandBuffer, PendingBarriers& pending);
  void recordPass(vk::CommandBuffer commandBuffer, Pass& pass);
  void recordRendering(vk::CommandBuffer commandBuffer, Pass& pass);

  std::vector<ImportedImage> images;
  std::vector<Pass> passes;
};

class RenderGraph::PassBuilder {
 public:
  /**
   * Renders into `image` as a color attachment. Passes with one or more of
   * these get rendering begun (and viewport/scissor set to the attachment's
   * extent) before their execute callback runs.
   */
  PassBuilder& writeColor(
      ImageRef image,
      vk::AttachmentLoadOp loadOp = vk::AttachmentLoadOp::eClear,
      ClearColor clearColor = {0.f, 0.f, 0.f, 1.f});

  /** Samples `image` in a shader. */
  PassBuilder& sample(ImageRef image);

  /** Any other use, e.g. storage or transfer access. */
  PassBuilder& use(ImageRef image, Access access);

 private:
  friend class RenderGraph;

  PassBuilder(RenderGraph& graph, Pass& pass) : graph{graph}, pass{pass} {}

  RenderGraph& graph;
  Pass& pass;
};

}  // namespace sophia