#include "graphics_pipeline.hpp"

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

namespace sophia {

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setVertexShader(
    Shader& vertexShader) {
  if (vertexShader.getStage() != ShaderStage::VERTEX) {
    log::fatal(
        "failed to configure graphics pipeline: shader is not a vertex "
        "shader");
    throw std::runtime_error(
        "failed to configure graphics pipeline: shader is not a vertex "
        "shader");
  }

  if (this->vertexShader != nullptr) {
    log::warning(
        "graphics pipeline builder: vertex shader already set, overwriting");
  }

  this->vertexShader = &vertexShader;
  return *this;
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setFragmentShader(
    Shader& fragmentShader) {
  if (fragmentShader.getStage() != ShaderStage::FRAGMENT) {
    log::fatal(
        "failed to configure graphics pipeline: shader is not a fragment "
        "shader");
    throw std::runtime_error(
        "failed to configure graphics pipeline: shader is not a fragment "
        "shader");
  }

  if (this->fragmentShader != nullptr) {
    log::warning(
        "graphics pipeline builder: fragment shader already set, "
        "overwriting");
  }

  this->fragmentShader = &fragmentShader;
  return *this;
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::addShader(
    Shader& shader) {
  switch (shader.getStage()) {
    case ShaderStage::VERTEX:
      return setVertexShader(shader);
    case ShaderStage::FRAGMENT:
      return setFragmentShader(shader);
    default:
      log::fatal(
          "failed to configure graphics pipeline: unsupported shader stage");
      throw std::runtime_error(
          "failed to configure graphics pipeline: unsupported shader stage");
  }
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setPipelineLayout(
    vk::PipelineLayout pipelineLayout) {
  this->pipelineLayout = pipelineLayout;
  return *this;
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setColorAttachmentFormats(
    std::vector<vk::Format> colorAttachmentFormats) {
  this->colorAttachmentFormats = std::move(colorAttachmentFormats);
  return *this;
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setDepthAttachmentFormat(
    vk::Format depthAttachmentFormat) {
  this->depthAttachmentFormat = depthAttachmentFormat;
  return *this;
}

GraphicsPipeline::Builder&
GraphicsPipeline::Builder::setStencilAttachmentFormat(
    vk::Format stencilAttachmentFormat) {
  this->stencilAttachmentFormat = stencilAttachmentFormat;
  return *this;
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setVertexInput(
    std::vector<vk::VertexInputBindingDescription> bindingDescriptions,
    std::vector<vk::VertexInputAttributeDescription> attributeDescriptions) {
  this->bindingDescriptions = std::move(bindingDescriptions);
  this->attributeDescriptions = std::move(attributeDescriptions);
  return *this;
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setAlphaBlending(
    bool enabled) {
  this->alphaBlending = enabled;
  return *this;
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setDepthTest(
    bool test, bool write, vk::CompareOp compareOp) {
  this->depthTest = test;
  this->depthWrite = write;
  this->depthCompareOp = compareOp;
  return *this;
}

GraphicsPipeline::Builder& GraphicsPipeline::Builder::setCullMode(
    vk::CullModeFlags cullMode, vk::FrontFace frontFace) {
  this->cullMode = cullMode;
  this->frontFace = frontFace;
  return *this;
}

std::unique_ptr<GraphicsPipeline> GraphicsPipeline::Builder::build() const {
  if (this->vertexShader == nullptr || this->fragmentShader == nullptr) {
    log::fatal(
        "failed to build graphics pipeline: setVertexShader()/"
        "setFragmentShader() was not called");
    throw std::runtime_error(
        "failed to build graphics pipeline: setVertexShader()/"
        "setFragmentShader() was not called");
  }

  if (this->pipelineLayout == VK_NULL_HANDLE) {
    log::fatal(
        "failed to build graphics pipeline: setPipelineLayout() was not "
        "called");
    throw std::runtime_error(
        "failed to build graphics pipeline: setPipelineLayout() was not "
        "called");
  }

  if (this->colorAttachmentFormats.empty() &&
      this->depthAttachmentFormat == vk::Format::eUndefined) {
    log::fatal(
        "failed to build graphics pipeline: no color or depth attachment "
        "format provided");
    throw std::runtime_error(
        "failed to build graphics pipeline: setColorAttachmentFormats() or "
        "setDepthAttachmentFormat() was not called");
  }

  return std::unique_ptr<GraphicsPipeline>(new GraphicsPipeline(*this));
}

GraphicsPipeline::GraphicsPipeline(const Builder& builder)
    : Pipeline(builder.device, vk::PipelineBindPoint::eGraphics) {
  setDefaultGraphicsPipelineConfig();

  if (builder.alphaBlending) {
    vk::PipelineColorBlendAttachmentState& blend =
        this->config.colorBlendAttachment;
    blend.blendEnable = vk::True;
    blend.srcColorBlendFactor = vk::BlendFactor::eSrcAlpha;
    blend.dstColorBlendFactor = vk::BlendFactor::eOneMinusSrcAlpha;
    blend.srcAlphaBlendFactor = vk::BlendFactor::eOne;
    blend.dstAlphaBlendFactor = vk::BlendFactor::eOneMinusSrcAlpha;
  }

  this->config.depthStencilInfo.depthTestEnable = builder.depthTest;
  this->config.depthStencilInfo.depthWriteEnable = builder.depthWrite;
  this->config.depthStencilInfo.depthCompareOp = builder.depthCompareOp;

  this->config.rasterizationInfo.cullMode = builder.cullMode;
  this->config.rasterizationInfo.frontFace = builder.frontFace;

  this->config.bindingDescriptions = builder.bindingDescriptions;
  this->config.attributeDescriptions = builder.attributeDescriptions;
  this->config.pipelineLayout = builder.pipelineLayout;
  this->config.colorAttachmentFormats = builder.colorAttachmentFormats;
  this->config.depthAttachmentFormat = builder.depthAttachmentFormat;
  this->config.stencilAttachmentFormat = builder.stencilAttachmentFormat;

  vk::UniqueShaderModule vertexShaderModule =
      createShaderModule(builder.vertexShader->getSpirv());
  vk::UniqueShaderModule fragmentShaderModule =
      createShaderModule(builder.fragmentShader->getSpirv());
  log::trace("created graphics shader modules");

  vk::PipelineShaderStageCreateInfo shaderStages[] = {
      {vk::PipelineShaderStageCreateFlags(), vk::ShaderStageFlagBits::eVertex,
       vertexShaderModule.get(), "main"},
      {vk::PipelineShaderStageCreateFlags(), vk::ShaderStageFlagBits::eFragment,
       fragmentShaderModule.get(), "main"}};

  vk::PipelineVertexInputStateCreateInfo vertexInputInfo{};
  vertexInputInfo.vertexBindingDescriptionCount =
      static_cast<u32>(this->config.bindingDescriptions.size());
  vertexInputInfo.pVertexBindingDescriptions =
      this->config.bindingDescriptions.data();
  vertexInputInfo.vertexAttributeDescriptionCount =
      static_cast<u32>(this->config.attributeDescriptions.size());
  vertexInputInfo.pVertexAttributeDescriptions =
      this->config.attributeDescriptions.data();

  vk::PipelineViewportStateCreateInfo viewportInfo{};
  viewportInfo.viewportCount = 1;
  viewportInfo.pViewports = nullptr;
  viewportInfo.scissorCount = 1;
  viewportInfo.pScissors = nullptr;

  vk::PipelineRenderingCreateInfo pipelineRenderingInfo{};
  pipelineRenderingInfo.colorAttachmentCount =
      static_cast<u32>(this->config.colorAttachmentFormats.size());
  pipelineRenderingInfo.pColorAttachmentFormats =
      this->config.colorAttachmentFormats.data();
  pipelineRenderingInfo.depthAttachmentFormat =
      this->config.depthAttachmentFormat;
  pipelineRenderingInfo.stencilAttachmentFormat =
      this->config.stencilAttachmentFormat;

  vk::GraphicsPipelineCreateInfo pipelineInfo{};
  pipelineInfo.pNext = &pipelineRenderingInfo;
  pipelineInfo.stageCount = 2;
  pipelineInfo.pStages = shaderStages;
  pipelineInfo.pVertexInputState = &vertexInputInfo;
  pipelineInfo.pInputAssemblyState = &this->config.inputAssemblyInfo;
  pipelineInfo.pViewportState = &viewportInfo;
  pipelineInfo.pRasterizationState = &this->config.rasterizationInfo;
  pipelineInfo.pMultisampleState = &this->config.multisampleInfo;
  pipelineInfo.pDepthStencilState = &this->config.depthStencilInfo;
  pipelineInfo.pColorBlendState = &this->config.colorBlendInfo;
  pipelineInfo.pDynamicState = &this->config.dynamicStateInfo;

  pipelineInfo.layout = this->config.pipelineLayout;
  pipelineInfo.basePipelineHandle = nullptr;

  try {
    this->pipeline = this->device.get()
                         ->createGraphicsPipelineUnique(nullptr, pipelineInfo)
                         .value;
    log::trace("created vk::Pipeline (graphics)");
  } catch (const vk::SystemError& err) {
    log::fatal("failed to create vk::Pipeline (graphics)");
    throw std::runtime_error("failed to create vk::Pipeline (graphics)");
  }
}

GraphicsPipeline::~GraphicsPipeline() {
  log::trace("destroyed vk::Pipeline (graphics)");
}

void GraphicsPipeline::setDefaultGraphicsPipelineConfig() {
  this->config.inputAssemblyInfo.topology =
      vk::PrimitiveTopology::eTriangleList;
  this->config.inputAssemblyInfo.primitiveRestartEnable = vk::False;

  this->config.rasterizationInfo.depthClampEnable = vk::False;
  this->config.rasterizationInfo.rasterizerDiscardEnable = vk::False;
  this->config.rasterizationInfo.polygonMode = vk::PolygonMode::eFill;
  this->config.rasterizationInfo.lineWidth = 1.0f;
  this->config.rasterizationInfo.cullMode = vk::CullModeFlagBits::eNone;
  this->config.rasterizationInfo.frontFace = vk::FrontFace::eCounterClockwise;
  this->config.rasterizationInfo.depthBiasEnable = vk::False;
  this->config.rasterizationInfo.depthBiasConstantFactor = 0.0f;
  this->config.rasterizationInfo.depthBiasClamp = 0.0f;
  this->config.rasterizationInfo.depthBiasSlopeFactor = 0.0f;

  this->config.multisampleInfo.sampleShadingEnable = vk::False;
  this->config.multisampleInfo.rasterizationSamples =
      vk::SampleCountFlagBits::e1;
  this->config.multisampleInfo.minSampleShading = 1.0f;
  this->config.multisampleInfo.pSampleMask = nullptr;
  this->config.multisampleInfo.alphaToCoverageEnable = vk::False;
  this->config.multisampleInfo.alphaToOneEnable = vk::False;

  this->config.colorBlendAttachment.colorWriteMask =
      vk::ColorComponentFlagBits::eR | vk::ColorComponentFlagBits::eG |
      vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eA;
  this->config.colorBlendAttachment.blendEnable = vk::False;
  this->config.colorBlendAttachment.srcColorBlendFactor = vk::BlendFactor::eOne;
  this->config.colorBlendAttachment.dstColorBlendFactor =
      vk::BlendFactor::eZero;
  this->config.colorBlendAttachment.colorBlendOp = vk::BlendOp::eAdd;
  this->config.colorBlendAttachment.srcAlphaBlendFactor = vk::BlendFactor::eOne;
  this->config.colorBlendAttachment.dstAlphaBlendFactor =
      vk::BlendFactor::eZero;
  this->config.colorBlendAttachment.alphaBlendOp = vk::BlendOp::eAdd;

  this->config.colorBlendInfo.logicOpEnable = vk::False;
  this->config.colorBlendInfo.logicOp = vk::LogicOp::eCopy;
  this->config.colorBlendInfo.attachmentCount = 1;
  this->config.colorBlendInfo.pAttachments = &this->config.colorBlendAttachment;
  this->config.colorBlendInfo.blendConstants[0] = 0.0f;
  this->config.colorBlendInfo.blendConstants[1] = 0.0f;
  this->config.colorBlendInfo.blendConstants[2] = 0.0f;
  this->config.colorBlendInfo.blendConstants[3] = 0.0f;

  this->config.depthStencilInfo.depthTestEnable = vk::True;
  this->config.depthStencilInfo.depthWriteEnable = vk::True;
  this->config.depthStencilInfo.depthCompareOp = vk::CompareOp::eLess;
  this->config.depthStencilInfo.depthBoundsTestEnable = vk::False;
  this->config.depthStencilInfo.minDepthBounds = 0.0f;
  this->config.depthStencilInfo.maxDepthBounds = 1.0f;
  this->config.depthStencilInfo.stencilTestEnable = vk::False;

  this->config.dynamicStateEnables = {vk::DynamicState::eViewport,
                                      vk::DynamicState::eScissor};
  this->config.dynamicStateInfo.pDynamicStates =
      this->config.dynamicStateEnables.data();
  this->config.dynamicStateInfo.dynamicStateCount =
      static_cast<u32>(this->config.dynamicStateEnables.size());
}

}  // namespace sophia