/*
	GameEngine and Editor
	by John Ryland
	Copyright (c) 2023
*/

////////////////////////////////////////////////////////////////////////////////////
//	Render View

// 1459
// 400

#include "RenderView.h"
#include "Image.h"
#include "imgui.h"
#include "VulkanRenderer.h"
#include "Vulkan.h"
#include "Model.h"

////////////////////////////////////////////////////////////////////////////////////
// Begin Rendering Code

#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#define GLM_ENABLE_EXPERIMENTAL
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtx/hash.hpp>

#define TINYOBJLOADER_IMPLEMENTATION
#include <tiny_obj_loader.h>

#include <cstddef>
#include <array>
#include <chrono>
#include <iostream>
#include <limits>
#include <stdexcept>
#include <unordered_map>
#include <vector>

const uint32_t WIDTH  = 800;
const uint32_t HEIGHT = 600;

// #define RESOURCES_PATH        "../vulkan/"
//#define MODEL_PATH            "models/viking_room.obj"
//#define TEXTURE_PATH          "textures/viking_room.png"

#define RESOURCES_PATH        "../GameEngine/resources/"
#define MODEL_PATH            "models/cube.obj"
#define TEXTURE_PATH          "textures/cube.png"
#define VERTEX_SHADER_PATH    "shaders/vert.spv"
#define FRAGMENT_SHADER_PATH  "shaders/frag.spv"

const int MAX_FRAMES_IN_FLIGHT = 3;

//struct Mat4 { alignas(16) glm::highp_f64mat4 value; };
struct Mat4 { alignas(16) glm::mat4 value; };

struct UniformBufferObject
{
    Mat4 m_model;
    Mat4 m_view;
    Mat4 m_projection;
};

template <size_t SIZE>
struct InFlightFrameData
{
    Vulkan::Buffer        m_uniformBuffers[SIZE];
    UniformBufferObject*  m_mappedUniformBuffers[SIZE];
    VkDescriptorSet       m_descriptorSets[SIZE];
    VkCommandBuffer       m_commandBuffers[SIZE];
    VkFence               m_fences[SIZE];
};

class ExampleRenderer
{
public:
    void Initialize(Vulkan::Device* device)
    {
        m_device = device;

        // We are rendering to an offscreen texture so we need to say that we want to also sample the color buffer after we have rendered in to it
        VkImageUsageFlags colorBufferUsage = /* VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT | */ VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
        VkFormat candidateDepthFormats[] = { VK_FORMAT_D32_SFLOAT, VK_FORMAT_D32_SFLOAT_S8_UINT, VK_FORMAT_D24_UNORM_S8_UINT };
        VkFormat depthFormat = m_device->FindFormat(candidateDepthFormats, std::size(candidateDepthFormats), VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT);
        uint32_t width = 512;
        uint32_t height = 512;
        Vulkan::Shader vertShader(m_device, RESOURCES_PATH VERTEX_SHADER_PATH);
        Vulkan::Shader fragShader(m_device, RESOURCES_PATH FRAGMENT_SHADER_PATH);

        m_model.Load(RESOURCES_PATH MODEL_PATH);
        m_renderPass.Initialize(m_device, VK_FORMAT_R8G8B8A8_UNORM, depthFormat, VK_SAMPLE_COUNT_1_BIT);
        m_pipeline.m_pipelineFlags = 0;
        m_pipeline.m_descriptorSetLayout = &m_device->m_descriptorSetLayout;
        m_pipeline.m_renderPass = m_renderPass.m_renderPass;
        m_pipeline.m_subpass = 0;
        m_pipeline.m_vertexBindingDescriptions = m_model.GetVertexBindingDescriptions();
        m_pipeline.m_vertexAttributeDescriptions = m_model.GetVertexAttributeDescriptions();
        m_pipeline.CreateLayout(m_device);
        m_pipeline.Initialize(m_device, vertShader.m_module, fragShader.m_module, VK_SAMPLE_COUNT_1_BIT);
        m_colorBuffer.Initialize(m_device, width, height, colorBufferUsage);
        m_depthBuffer.Initialize(m_device, width, height, depthFormat);
        m_framebuffer.Initialize(m_device, m_renderPass.m_renderPass, &m_colorBuffer, &m_depthBuffer);
        m_texture.Initialize(m_device, RESOURCES_PATH TEXTURE_PATH);

        createVertexBuffer();
        createIndexBuffer();
        createCommandBuffers();
        createUniformBuffers();
        createDescriptorSets();
        createSyncObjects();
    }

    void Shutdown()
    {
        vkDeviceWaitIdle(m_device->m_device);

        // The color buffer is managed by the main Renderer now so don't delete it
        m_colorBuffer.m_view = nullptr;

        m_colorBuffer.Destroy();
        m_depthBuffer.Destroy();
        m_pipeline.Destroy();
        m_renderPass.Destroy();

        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
            inFlightData.m_uniformBuffers[i].Destroy();

        m_texture.Destroy();
        m_indexBuffer.Destroy();
        m_vertexBuffer.Destroy();

        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
            vkDestroyFence(m_device->m_device, inFlightData.m_fences[i], nullptr);
    }

    void Update()
    {
        vkWaitForFences(m_device->m_device, 1, &inFlightData.m_fences[currentFrame], VK_TRUE, UINT64_MAX);
        updateUniformBuffer();
        vkResetFences(m_device->m_device, 1, &inFlightData.m_fences[currentFrame]);
        vkResetCommandBuffer(inFlightData.m_commandBuffers[currentFrame], /*VkCommandBufferResetFlagBits*/ 0);
        recordCommandBuffer(inFlightData.m_commandBuffers[currentFrame]);

        VkSubmitInfo submitInfo {};
        submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
        submitInfo.commandBufferCount = 1;
        submitInfo.pCommandBuffers    = &inFlightData.m_commandBuffers[currentFrame];
        m_device->CheckResult(vkQueueSubmit(m_device->m_graphicsQueue, 1, &submitInfo, inFlightData.m_fences[currentFrame]), "failed to submit draw command buffer!");

        currentFrame = (currentFrame + 1) % MAX_FRAMES_IN_FLIGHT;
    }

    VkImageView ImageView()
    {
        return m_colorBuffer.m_view;
    }

private:
    // Context
    Vulkan::Device*              m_device;       // TODO: rename Device -> Context
    Vulkan::Pipeline             m_pipeline;
    Vulkan::RenderPass           m_renderPass;

    // Framebuffer
    Vulkan::ImageBuffer          m_colorBuffer;
    Vulkan::DepthBuffer          m_depthBuffer;
    Vulkan::Framebuffer          m_framebuffer;

    // Scene
    Vulkan::Texture              m_texture;
    Vulkan::Buffer               m_vertexBuffer;
    Vulkan::Buffer               m_indexBuffer;

    // Data
    Render::Model                m_model;

    // Original - per swap-chain frame data
    uint32_t                                 currentFrame = 0;
    InFlightFrameData<MAX_FRAMES_IN_FLIGHT>  inFlightData;

    void createVertexBuffer()
    {
        VkDeviceSize bufferSize;
        const uint8_t* data;
        m_model.GetVertexData(data, bufferSize);
        m_vertexBuffer.Initialize(m_device, bufferSize, VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        m_vertexBuffer.UploadDataUsingStagingBuffer(data);
    }

    void createIndexBuffer()
    {
        VkDeviceSize bufferSize;
        const uint8_t* data;
        m_model.GetIndexData(data, bufferSize);
        m_indexBuffer.Initialize(m_device, bufferSize, VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
        m_indexBuffer.UploadDataUsingStagingBuffer(data);
    }

    void createCommandBuffers()
    {
        VkCommandBufferAllocateInfo allocInfo {};
        allocInfo.sType              = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        allocInfo.commandPool        = m_device->m_commandPool;
        allocInfo.level              = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        allocInfo.commandBufferCount = std::size(inFlightData.m_commandBuffers);
        m_device->CheckResult(vkAllocateCommandBuffers(m_device->m_device, &allocInfo, inFlightData.m_commandBuffers), "failed to allocate command buffers!");
    }

    void createUniformBuffers()
    {
        VkDeviceSize bufferSize = sizeof(UniformBufferObject);
        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
        {
            inFlightData.m_uniformBuffers[i].Initialize(m_device, bufferSize, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT);
            vkMapMemory(m_device->m_device, inFlightData.m_uniformBuffers[i].m_memory, 0, bufferSize, 0, (void**)&inFlightData.m_mappedUniformBuffers[i]);
        }
    }

    void createSyncObjects()
    {
        VkFenceCreateInfo fenceInfo {};
        fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
        fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;
        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
            m_device->CheckResult(vkCreateFence(m_device->m_device, &fenceInfo, nullptr, &inFlightData.m_fences[i]), "failed to create synchronization objects for a frame!");
    }

    void createDescriptorSets()
    {
        VkDescriptorSetLayout layouts[MAX_FRAMES_IN_FLIGHT];
        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
            layouts[i] = m_device->m_descriptorSetLayout;

        VkDescriptorSetAllocateInfo allocInfo {};
        allocInfo.sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
        allocInfo.descriptorPool     = m_device->m_descriptorPool;
        allocInfo.descriptorSetCount = std::size(layouts);
        allocInfo.pSetLayouts        = layouts;

        m_device->CheckResult(vkAllocateDescriptorSets(m_device->m_device, &allocInfo, inFlightData.m_descriptorSets), "failed to allocate descriptor sets!");

        for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++)
        {
            VkDescriptorBufferInfo bufferInfo {};
            bufferInfo.buffer = inFlightData.m_uniformBuffers[i].m_buffer;
            bufferInfo.offset = 0;
            bufferInfo.range  = sizeof(UniformBufferObject);

            VkDescriptorImageInfo imageInfo {};
            imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
            imageInfo.imageView   = m_texture.m_view;
            imageInfo.sampler     = m_texture.m_sampler;

            std::array<VkWriteDescriptorSet, 2> descriptorWrites {};

            descriptorWrites[0].sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            descriptorWrites[0].dstSet          = inFlightData.m_descriptorSets[i];
            descriptorWrites[0].dstBinding      = 0;
            descriptorWrites[0].dstArrayElement = 0;
            descriptorWrites[0].descriptorType  = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
            descriptorWrites[0].descriptorCount = 1;
            descriptorWrites[0].pBufferInfo     = &bufferInfo;

            descriptorWrites[1].sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            descriptorWrites[1].dstSet          = inFlightData.m_descriptorSets[i];
            descriptorWrites[1].dstBinding      = 1;
            descriptorWrites[1].dstArrayElement = 0;
            descriptorWrites[1].descriptorType  = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
            descriptorWrites[1].descriptorCount = 1;
            descriptorWrites[1].pImageInfo      = &imageInfo;

            vkUpdateDescriptorSets(m_device->m_device,
                                   static_cast<uint32_t>(descriptorWrites.size()),
                                   descriptorWrites.data(),
                                   0,
                                   nullptr);
        }
    }

    void recordCommandBuffer(VkCommandBuffer commandBuffer)
    {
        VkCommandBufferBeginInfo beginInfo {};
        beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
        m_device->CheckResult(vkBeginCommandBuffer(commandBuffer, &beginInfo), "failed to begin recording command buffer!");

        std::array<VkClearValue, 2> clearValues {};
        clearValues[0].color = { { 0.0f, 0.0f, 0.0f, 1.0f } };
        clearValues[1].depthStencil = { 1.0f, 0 };

        VkRenderPassBeginInfo renderPassInfo {};
        renderPassInfo.sType             = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
        renderPassInfo.renderPass        = m_renderPass.m_renderPass;
        renderPassInfo.framebuffer       = m_framebuffer.m_framebuffer;
        renderPassInfo.renderArea.offset = { 0, 0 };
        renderPassInfo.renderArea.extent.width  = m_framebuffer.m_width;
        renderPassInfo.renderArea.extent.height = m_framebuffer.m_height;
        renderPassInfo.clearValueCount   = static_cast<uint32_t>(clearValues.size());
        renderPassInfo.pClearValues      = clearValues.data();
        vkCmdBeginRenderPass(commandBuffer, &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);
        vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_pipeline.m_pipeline);

        VkViewport viewport {};
        viewport.x        = 0.0f;
        viewport.y        = 0.0f;
        viewport.width    = (float)m_framebuffer.m_width;
        viewport.height   = (float)m_framebuffer.m_height;
        viewport.minDepth = 0.0f;
        viewport.maxDepth = 1.0f;
        vkCmdSetViewport(commandBuffer, 0, 1, &viewport);

        VkRect2D scissor {};
        scissor.offset = { 0, 0 };
        scissor.extent.width  = m_framebuffer.m_width;
        scissor.extent.height = m_framebuffer.m_height;
        vkCmdSetScissor(commandBuffer, 0, 1, &scissor);

        VkBuffer     vertexBuffers[] = { m_vertexBuffer.m_buffer };
        VkDeviceSize offsets[]       = { 0 };
        vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
        vkCmdBindIndexBuffer(commandBuffer, m_indexBuffer.m_buffer, 0, VK_INDEX_TYPE_UINT32);
        vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_pipeline.m_pipelineLayout, 0, 1, &inFlightData.m_descriptorSets[currentFrame], 0, nullptr);
        vkCmdDrawIndexed(commandBuffer, m_model.GetIndexCount(), 1, 0, 0, 0);
        vkCmdEndRenderPass(commandBuffer);
        m_device->CheckResult(vkEndCommandBuffer(commandBuffer), "failed to record command buffer!");
    }

    void updateUniformBuffer()
    {
        static auto startTime = std::chrono::high_resolution_clock::now();
        auto  currentTime = std::chrono::high_resolution_clock::now();

        // TODO : float is bad for this
        float time        = std::chrono::duration<float, std::chrono::seconds::period>(currentTime - startTime).count();

        UniformBufferObject ubo {};
        ubo.m_model.value      = glm::rotate(glm::mat4(1.0f), time * glm::radians(90.0f), glm::vec3(0.0f, 0.0f, 1.0f));
        ubo.m_view.value       = glm::lookAt(glm::vec3(2.0f, 2.0f, 2.0f), glm::vec3(0.0f, 0.0f, 0.0f), glm::vec3(0.0f, 0.0f, 1.0f));
        ubo.m_projection.value = glm::perspective(glm::radians(45.0f),
                                            m_framebuffer.m_width / (float)m_framebuffer.m_height,
                                            0.1f,
                                            10.0f);
        ubo.m_projection.value[1][1] *= -1;
        memcpy(inFlightData.m_mappedUniformBuffers[currentFrame], &ubo, sizeof(ubo));
    }
};

// End Rendering Code
////////////////////////////////////////////////////////////////////////////////////



namespace GameEngine {

RenderView::RenderView(ApplicationFramework::IApplication& app)
    : m_renderer(new ExampleRenderer())
    , m_application(app)
{
}

RenderView::~RenderView()
{
}

void RenderView::AddShowMenuItem()
{
    ImGui::MenuItem("Show Render View", NULL, &m_open);
}

void RenderView::Update()
{
    if (m_open)
    {
        if (ImGui::Begin("Render View", &m_open))
        {
            ImGui::Text("pointer = %llu", (ImTextureID)m_texture);
            ImGui::Text("size = %d x %d", m_width, m_height);
            ImGui::Image((ImTextureID)m_texture, ImVec2(m_width, m_height));
            m_renderer->Update();
        }
        ImGui::End();
    }
}

void RenderView::Initialize()
{
    m_renderer->Initialize(((ApplicationFramework::VulkanRenderer*)m_application.RenderSystem())->Details());
    m_texture = m_application.RenderSystem()->LoadTexture(m_renderer->ImageView(), 512, 512);
    m_width = 512;
    m_height = 512;
}

void RenderView::Shutdown()
{
    m_application.RenderSystem()->DestroyTexture(m_texture);
    m_renderer->Shutdown();
    delete m_renderer;
    m_renderer = nullptr;
}

} // GameEngine namespace