#include "Buffer.h"
#include "DescriptorPool.h"
#include "DescriptorSet.h"
#include "Framebuffer.h"
#include "IndexBuffer.h"
#include "Instance.h"
#include "Pipeline.h"
#include "Texture2D.h"
#include "Timer.h"
#include "UniformBuffer.h"
#include "Vertex.h"
#include "VertexBuffer.h"
#include "VertexPassthrough.h"

#include <GLFW/glfw3.h>
#include <chrono>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <iostream>
#include <optional>
#include <set>
#include <vector>

const std::vector<Vertex> vertices = {
  Vertex{{-0.5f, 0.5f, -0.5f}, {1.0f, 0.0f, 0.0f}, {0.0f, 0.0f}},
  Vertex{{ 0.5f, 0.5f, -0.5f}, {0.0f, 1.0f, 0.0f}, {1.0f, 0.0f}},
  Vertex{{ 0.5f, 0.5f,  0.5f}, {0.0f, 0.0f, 1.0f}, {1.0f, 1.0f}},
  Vertex{{-0.5f, 0.5f,  0.5f}, {1.0f, 1.0f, 1.0f}, {0.0f, 1.0f}},
  Vertex{{-0.5f, 0.0f, -0.5f}, {1.0f, 0.0f, 0.0f}, {0.0f, 0.0f}},
  Vertex{{ 0.5f, 0.0f, -0.5f}, {0.0f, 1.0f, 0.0f}, {1.0f, 0.0f}},
  Vertex{{ 0.5f, 0.0f,  0.5f}, {0.0f, 0.0f, 1.0f}, {1.0f, 1.0f}},
  Vertex{{-0.5f, 0.0f,  0.5f}, {1.0f, 1.0f, 1.0f}, {0.0f, 1.0f}},
};

const std::vector<uint16_t> indices = {
    0, 1, 2, 2, 3, 0,
    4, 5, 6, 6, 7, 4
};

const std::vector<VertexPassthrough> verticesPassthrough = {
  VertexPassthrough{{-1.0f, -1.0f}},
  VertexPassthrough{{ 1.0f, -1.0f}},
  VertexPassthrough{{ 1.0f,  1.0f}},
  VertexPassthrough{{-1.0f,  1.0f}},
};

const std::vector<uint16_t> indicesPassthrough = {
    0, 1, 2, 2, 3, 0,
};

struct alignas(64) ShaderUniform
{
  alignas(16) glm::mat4 projection;
  alignas(16) glm::mat4 view;
  alignas(16) glm::mat4 model;
  alignas(16) glm::vec3 lightPos;
};

class Application final
{
private:
  std::unique_ptr<Instance> instance;
  std::unique_ptr<Pipeline> graphicsPipeline;
  std::unique_ptr<Texture2D> texture2D;
  std::unique_ptr<UniformBuffer> shaderUniformBuffer;
  std::unique_ptr<DescriptorPool> descriptorPool;
  std::unique_ptr<DescriptorSet> descriptorSet;
  std::unique_ptr<VertexBuffer> vertexBuffer;
  std::unique_ptr<IndexBuffer> indexBuffer;
  std::unique_ptr<CommandBuffer> commandBuffer;

  std::unique_ptr<Framebuffer> framebuffer;
  std::unique_ptr<Pipeline> graphicsPipelinePassthrough;
  std::unique_ptr<VertexBuffer> vertexBufferPassthrough;
  std::unique_ptr<IndexBuffer> indexBufferPassthrough;
  std::unique_ptr<DescriptorSet> descriptorSetPassthrough;

public:
  Application()
  {
    InitializeInstance();
    InitializeFrameBuffer();
    InitializeGraphicsPipeline();
    InitializeTextureSampler();
    InitializeUniformBuffer();
    InitializeDescriptorSets();
    InitializeVertexBuffer();
    InitializeIndexBuffer();
    InitializeCommandBuffer();
  }

  ~Application()
  {
    vkDeviceWaitIdle(instance->GetDevice());
  }

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

  bool Update()
  {
    if (!instance->BeginPresent())
      return true;

    RecordCommandBuffer();
    commandBuffer->SubmitAsGraphicsQueue();

    return instance->EndPresent();
  }

private:

  void InitializeInstance()
  {
    instance = std::make_unique<Instance>("Copium Engine");
  }

  void InitializeFrameBuffer()
  {
    framebuffer = std::make_unique<Framebuffer>(*instance, instance->GetSwapChain().GetExtent().width, instance->GetSwapChain().GetExtent().height);
  }

  void InitializeTextureSampler()
  {
    texture2D = std::make_unique<Texture2D>(*instance, "res/textures/texture.png");
  }

  void InitializeUniformBuffer()
  {
    shaderUniformBuffer = std::make_unique<UniformBuffer>(*instance, sizeof(ShaderUniform));
  }

  void InitializeDescriptorSets()
  {
    descriptorPool = std::make_unique<DescriptorPool>(*instance);

    descriptorSet = std::make_unique<DescriptorSet>(*instance, *descriptorPool, graphicsPipeline->GetDescriptorSetLayout(0));
    descriptorSet->AddUniform(*shaderUniformBuffer, 0);
    descriptorSet->AddTexture2D(*texture2D, 1);

    descriptorSetPassthrough = std::make_unique<DescriptorSet>(*instance, *descriptorPool, graphicsPipelinePassthrough->GetDescriptorSetLayout(0));
    descriptorSetPassthrough->AddTexture2D(framebuffer->GetTexture2D(), 0);
  }

  void InitializeGraphicsPipeline() 
  {
    PipelineCreator creator{framebuffer->GetRenderPass(), "res/shaders/shader.vert", "res/shaders/shader.frag"};
    creator.AddDescriptorSetLayoutBinding(0, 0, VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1, VK_SHADER_STAGE_VERTEX_BIT);
    creator.AddDescriptorSetLayoutBinding(0, 1, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1, VK_SHADER_STAGE_FRAGMENT_BIT);
    creator.SetVertexDescriptor(Vertex::GetDescriptor());
    creator.SetCullMode(VK_CULL_MODE_NONE);
    graphicsPipeline = std::make_unique<Pipeline>(*instance, creator);

    PipelineCreator creatorPassthrough{instance->GetSwapChain().GetRenderPass(), "res/shaders/passthrough.vert", "res/shaders/passthrough.frag"};
    creatorPassthrough.AddDescriptorSetLayoutBinding(0, 0, VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1, VK_SHADER_STAGE_FRAGMENT_BIT);
    creatorPassthrough.SetVertexDescriptor(VertexPassthrough::GetDescriptor());
    creatorPassthrough.SetCullMode(VK_CULL_MODE_NONE);
    graphicsPipelinePassthrough = std::make_unique<Pipeline>(*instance, creatorPassthrough);
  }

  void InitializeVertexBuffer()
  {
    vertexBuffer = std::make_unique<VertexBuffer>(*instance, Vertex::GetDescriptor(), vertices.size());
    vertexBuffer->Update(0, (void*)vertices.data());

    vertexBufferPassthrough = std::make_unique<VertexBuffer>(*instance, VertexPassthrough::GetDescriptor(), verticesPassthrough.size());
    vertexBufferPassthrough->Update(0, (void*)verticesPassthrough.data());
  }

  void InitializeIndexBuffer()
  {
    indexBuffer = std::make_unique<IndexBuffer>(*instance, indices.size());
    indexBuffer->UpdateStaging((void*)indices.data());

    indexBufferPassthrough = std::make_unique<IndexBuffer>(*instance, indicesPassthrough.size());
    indexBufferPassthrough->UpdateStaging((void*)indicesPassthrough.data());
  }
  
  void InitializeCommandBuffer()
  {
    commandBuffer = std::make_unique<CommandBuffer>(*instance, CommandBufferType::Dynamic);
  }

  void RecordCommandBuffer()
  {
    commandBuffer->Begin();
    std::vector<VkClearValue> clearValues{2};
    clearValues[0].color = {{0.0f, 0.0f, 0.0f, 1.0f}};
    clearValues[1].depthStencil = {1.0f, 0};



    framebuffer->Bind(*commandBuffer);
    graphicsPipeline->Bind(*commandBuffer);

    UpdateUniformBuffer();

    vertexBuffer->Bind(*commandBuffer);
    indexBuffer->Bind(*commandBuffer);

    graphicsPipeline->SetDescriptorSet(0, *descriptorSet);
    graphicsPipeline->BindDescriptorSets(commandBuffer->GetHandle());

    indexBuffer->Draw(*commandBuffer);
    framebuffer->Unbind(*commandBuffer);

    instance->GetSwapChain().BeginFrameBuffer(*commandBuffer);

    graphicsPipelinePassthrough->Bind(*commandBuffer);
    graphicsPipelinePassthrough->SetDescriptorSet(0, *descriptorSetPassthrough);
    graphicsPipelinePassthrough->BindDescriptorSets(commandBuffer->GetHandle());
    vertexBufferPassthrough->Bind(*commandBuffer);
    indexBufferPassthrough->Bind(*commandBuffer);
    indexBufferPassthrough->Draw(*commandBuffer);

    instance->GetSwapChain().EndFrameBuffer(*commandBuffer);
    commandBuffer->End();
  }

  void UpdateUniformBuffer()
  {
    static Timer startTimer;

    float time = startTimer.Elapsed();
    ShaderUniform shaderUniform;
    shaderUniform.view = glm::lookAt(glm::vec3(2.0f, 2.0f, 2.0f), glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f, 1.0f, 0.0f));
    shaderUniform.projection = glm::perspective(glm::radians(45.0f), instance->GetSwapChain().GetExtent().width / (float)instance->GetSwapChain().GetExtent().height, 0.1f, 10.0f);
    shaderUniform.model = glm::rotate(glm::mat4(1.0f), time * glm::radians(90.0f), glm::vec3(0.0f, 1.0f, 0.0f));
    shaderUniform.projection[1][1] *= -1;
    shaderUniform.lightPos = glm::rotate(glm::mat4{1.0f}, time * glm::radians(45.0f), glm::vec3(0, 1, 0)) * glm::vec4{0.3, 0.1, 0, 1};

    shaderUniformBuffer->Update(shaderUniform);
  }
};

void func(const int* ptr) {
  *const_cast<int*>(ptr) = 20;
}

int main()
{
  CP_ASSERT(glfwInit() == GLFW_TRUE, "main : Failed to initialize the glfw context");
  {
    Application application;
    Timer timer;
    int frames = 0;
    while (application.Update())
    {
      glfwPollEvents();
      if (timer.Elapsed() >= 1.0)
      {
        CP_DEBUG("%d fps", frames);
        frames = 0;
        timer.Start();
      }
      frames++;
    }
  }

  glfwTerminate();
  return 0;
}