#include "UwURenderEngine.hpp"

#include <array>
#include <chrono>

#include <glm/gtc/matrix_transform.hpp>

#include "GPU_GarbageCollector.h"
#include "../Game/Actors/Camera.hpp"
#include "Core/CoreInstance.hpp"

#include "Game/GameInstance.hpp"
#include "Game/World/World.hpp"
#include "../Game/Actors/Mesh/Mesh.hpp"
#include "DeviceFunctions/DeviceFunctions.hpp"

RENDER_ENGINE_FACTORY_GENERATE(UwURenderEngine)

void UwURenderEngine::create_uniform_buffers()
{
    uniform_buffers_.resize(swapchain_.get_images().size());
    uniform_buffers_memory_.resize(swapchain_.get_images().size());
    uniform_buffers_mapped_.resize(swapchain_.get_images().size());

    std::vector<uint32_t> arr_indices = physical_device_.get_unique_family_indices();
    for (size_t i = 0; i < swapchain_.get_images().size(); ++i)
    {
        constexpr VkDeviceSize size = sizeof(UniformBufferObject);
        uniform_buffers_[i] = DeviceFunctions::create_buffer(device_.get_native(), size, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, arr_indices);
        uniform_buffers_memory_[i] = DeviceFunctions::allocate_device_memory(physical_device_.get_native(), device_.get_native(), uniform_buffers_[i], VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
        vkBindBufferMemory(device_.get_native(), uniform_buffers_[i], uniform_buffers_memory_[i], 0);
    
        vkMapMemory(device_.get_native(), uniform_buffers_memory_[i], 0, size, 0, &uniform_buffers_mapped_[i]);
    }
}

void UwURenderEngine::create_descriptor_pool()
{
    VkDescriptorPoolSize pool_size;
    pool_size.type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    pool_size.descriptorCount = static_cast<uint32_t>(swapchain_.get_images().size());

    VkDescriptorPoolCreateInfo descriptor_pool_info{};
    descriptor_pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    descriptor_pool_info.poolSizeCount = 1;
    descriptor_pool_info.pPoolSizes = &pool_size;
    descriptor_pool_info.maxSets = static_cast<uint32_t>(swapchain_.get_images().size());

    VK_CHECK(vkCreateDescriptorPool(device_.get_native(), &descriptor_pool_info, nullptr, &descriptor_pool_))
}

void UwURenderEngine::create_descriptor_sets()
{
    std::vector<VkDescriptorSetLayout> layouts(swapchain_.get_images().size(), layout_binding_.get_native());
    VkDescriptorSetAllocateInfo descriptor_set_allocate_info{};
    descriptor_set_allocate_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    descriptor_set_allocate_info.descriptorPool = descriptor_pool_;
    descriptor_set_allocate_info.descriptorSetCount = static_cast<uint32_t>(swapchain_.get_images().size());
    descriptor_set_allocate_info.pSetLayouts = layouts.data();

    descriptor_sets_.resize(swapchain_.get_images().size());
    VK_CHECK(vkAllocateDescriptorSets(device_.get_native(), &descriptor_set_allocate_info, descriptor_sets_.data()))

    for (size_t i = 0; i < swapchain_.get_images().size(); ++i)
    {
        VkDescriptorBufferInfo buffer_info{};
        buffer_info.buffer = uniform_buffers_[i];
        buffer_info.offset = 0;
        buffer_info.range = sizeof(UniformBufferObject);

        VkWriteDescriptorSet write_descriptor_set{};
        write_descriptor_set.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        write_descriptor_set.dstSet = descriptor_sets_[i];
        write_descriptor_set.dstBinding = 0;
        write_descriptor_set.dstArrayElement = 0;
        write_descriptor_set.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        write_descriptor_set.descriptorCount = 1;
        write_descriptor_set.pBufferInfo = &buffer_info;
        write_descriptor_set.pImageInfo = nullptr;
        write_descriptor_set.pTexelBufferView = nullptr;

        vkUpdateDescriptorSets(device_.get_native(), 1, &write_descriptor_set, 0, nullptr);
    }
}

void UwURenderEngine::create_framebuffers()
{
    framebuffers_.resize(swapchain_.get_images().size());
    
    VkFramebufferCreateInfo framebuffer_info{};
    framebuffer_info.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
    framebuffer_info.renderPass = render_pass_.get_native();
    framebuffer_info.attachmentCount = 1;
    framebuffer_info.width = swapchain_.get_extent().width;
    framebuffer_info.height = swapchain_.get_extent().height;
    framebuffer_info.layers = 1;

    for (size_t i = 0; i < framebuffers_.size(); ++i)
    {
        framebuffer_info.pAttachments = &swapchain_.get_image_views()[i];
        VK_CHECK(vkCreateFramebuffer(device_.get_native(), &framebuffer_info, nullptr, &framebuffers_[i]))
    }
}

void UwURenderEngine::create_command_pool()
{
    VkObjects::PhysicalDevice::QueueFamilyIndices queue_family_indices = physical_device_.get_queue_family_indices();
    
    VkCommandPoolCreateInfo command_pool_info{};
    command_pool_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    command_pool_info.queueFamilyIndex = queue_family_indices.graphics_family;
    command_pool_info.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;

    VK_CHECK(vkCreateCommandPool(device_.get_native(), &command_pool_info, nullptr, &command_pool_))
}

void UwURenderEngine::allocate_vertex_buffer()
{
    VkObjects::PhysicalDevice::QueueFamilyIndices indices = physical_device_.get_queue_family_indices();
    
    
    std::vector<uint32_t> arr_indices = physical_device_.get_unique_family_indices();
    std::vector<uint32_t> transfer_index = {indices.transfer_family};

    VkDeviceSize size_buffer = sizeof(vertices_[0]) * vertices_.size();
    VkBuffer staging_buffer = DeviceFunctions::create_buffer(device_.get_native(), size_buffer, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                                                             transfer_index);
    VkDeviceMemory staging_buffer_memory = DeviceFunctions::allocate_device_memory(
        physical_device_.get_native(), device_.get_native(), staging_buffer,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
    
    vertex_buffer_ = DeviceFunctions::create_buffer(device_.get_native(), size_buffer, VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT, arr_indices);
    vertex_buffer_memory_ = DeviceFunctions::allocate_device_memory(physical_device_.get_native(), device_.get_native(), vertex_buffer_, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
    
    VK_CHECK(vkBindBufferMemory(device_.get_native(), staging_buffer, staging_buffer_memory, 0))
    VK_CHECK(vkBindBufferMemory(device_.get_native(), vertex_buffer_, vertex_buffer_memory_, 0))

    void* data;
    VK_CHECK(vkMapMemory(device_.get_native(), staging_buffer_memory, 0, size_buffer, 0, &data))
    memcpy(data, vertices_.data(), size_buffer);
    vkUnmapMemory(device_.get_native(), staging_buffer_memory);

    DeviceFunctions::device_memcpy(device_.get_native(), indices.transfer_family, staging_buffer, vertex_buffer_, size_buffer);

    vkFreeMemory(device_.get_native(), staging_buffer_memory, nullptr);
    vkDestroyBuffer(device_.get_native(), staging_buffer, nullptr);
}

void UwURenderEngine::allocate_index_buffer()
{
    std::vector<uint32_t> arr_indices = physical_device_.get_unique_family_indices();
    VkObjects::PhysicalDevice::QueueFamilyIndices indices = physical_device_.get_queue_family_indices();
    std::vector<uint32_t> transfer_index = {indices.transfer_family};

    VkDeviceSize size_buffer = sizeof(indices_[0]) * indices_.size();
    VkBuffer staging_buffer = DeviceFunctions::create_buffer(device_.get_native(), size_buffer, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                                                             transfer_index);
    VkDeviceMemory staging_buffer_memory = DeviceFunctions::allocate_device_memory(
        physical_device_.get_native(), device_.get_native(), staging_buffer,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
    vkBindBufferMemory(device_.get_native(), staging_buffer, staging_buffer_memory, 0);
    
    void* data;
    vkMapMemory(device_.get_native(), staging_buffer_memory, 0, size_buffer, 0, &data);
    memcpy(data, indices_.data(), size_buffer);
    vkUnmapMemory(device_.get_native(), staging_buffer_memory);
    
    index_buffer_ = DeviceFunctions::create_buffer(device_.get_native(), size_buffer, VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT, arr_indices);
    index_buffer_memory_ = DeviceFunctions::allocate_device_memory(physical_device_.get_native(), device_.get_native(), index_buffer_, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
    vkBindBufferMemory(device_.get_native(), index_buffer_, index_buffer_memory_, 0);

    DeviceFunctions::device_memcpy(device_.get_native(), indices.transfer_family,staging_buffer, index_buffer_, size_buffer);

    vkDestroyBuffer(device_.get_native(), staging_buffer, nullptr);
    vkFreeMemory(device_.get_native(), staging_buffer_memory, nullptr);
}

void UwURenderEngine::allocate_command_buffers()
{
    command_buffers_.resize(swapchain_.get_images().size());

    VkCommandBufferAllocateInfo command_buffer_allocate_info{};
    command_buffer_allocate_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    command_buffer_allocate_info.commandPool = command_pool_;
    command_buffer_allocate_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    command_buffer_allocate_info.commandBufferCount = static_cast<uint32_t>(command_buffers_.size());

    VK_CHECK(vkAllocateCommandBuffers(device_.get_native(), &command_buffer_allocate_info, command_buffers_.data()))
}

void UwURenderEngine::create_sync_objects()
{
    image_available_semaphores_.resize(swapchain_.get_images().size());
    render_finished_semaphores_.resize(swapchain_.get_images().size());
    in_flight_fences_.resize(swapchain_.get_images().size());

    VkSemaphoreCreateInfo semaphore_info{};
    semaphore_info.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;

    VkFenceCreateInfo fence_info{};
    fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
    fence_info.flags = VK_FENCE_CREATE_SIGNALED_BIT;

    for (size_t i = 0; i < swapchain_.get_images().size(); ++i)
    {
        VK_CHECK(vkCreateSemaphore(device_.get_native(), &semaphore_info, nullptr, &image_available_semaphores_[i]))
        VK_CHECK(vkCreateSemaphore(device_.get_native(), &semaphore_info, nullptr, &render_finished_semaphores_[i]))
        VK_CHECK(vkCreateFence(device_.get_native(), &fence_info, nullptr, &in_flight_fences_[i]))
    }
}

void UwURenderEngine::record_command_buffer(VkCommandBuffer command_buffer, uint32_t image_index) const
{
    VkCommandBufferBeginInfo command_buffer_begin_info{};
    command_buffer_begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    VK_CHECK(vkBeginCommandBuffer(command_buffer, &command_buffer_begin_info))

    constexpr VkClearValue clear_color = {{{0.0f, 0.0f, 0.0f, 1.0f}}};
    VkRenderPassBeginInfo render_pass_begin_info{};
    render_pass_begin_info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    render_pass_begin_info.renderPass = render_pass_.get_native();
    render_pass_begin_info.framebuffer = framebuffers_[image_index];
    render_pass_begin_info.renderArea.offset = {0, 0};
    render_pass_begin_info.renderArea.extent = swapchain_.get_extent();
    render_pass_begin_info.clearValueCount = 1;
    render_pass_begin_info.pClearValues = &clear_color;
    vkCmdBeginRenderPass(command_buffer, &render_pass_begin_info, VK_SUBPASS_CONTENTS_INLINE);

    vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_.get_native());

    VkBuffer vertex_buffers[] = {vertex_buffer_};
    VkDeviceSize offset[] = {0};
    vkCmdBindVertexBuffers(command_buffer, 0, 1, vertex_buffers, offset);

    vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_.get_layout(), 0, 1, &descriptor_sets_[current_frame_], 0, nullptr);
    
    vkCmdBindIndexBuffer(command_buffer, index_buffer_, 0, VK_INDEX_TYPE_UINT16);

    VkViewport viewport;
    viewport.x = 0.0f;
    viewport.y = 0.0f;
    viewport.width = static_cast<float>(swapchain_.get_extent().width);
    viewport.height = static_cast<float>(swapchain_.get_extent().height);
    viewport.minDepth = 0.0f;
    viewport.maxDepth = 1.0f;
    vkCmdSetViewport(command_buffer, 0, 1, &viewport);

    VkRect2D scissor;
    scissor.offset = {0, 0};
    scissor.extent = swapchain_.get_extent();
    vkCmdSetScissor(command_buffer, 0, 1, &scissor);

    vkCmdDrawIndexed(command_buffer, static_cast<uint32_t>(indices_.size()), 1, 0, 0, 0);
    
    vkCmdEndRenderPass(command_buffer);
    VK_CHECK(vkEndCommandBuffer(command_buffer))
}

void UwURenderEngine::update_uniform_buffer(uint32_t current_frame) const
{
    glm::vec3 eye{0, 0, 0};
    glm::vec3 center{0, 0, 0};
    glm::vec3 up{0, 1, 0};
    if (active_camera_.expired() == false)
    {
        std::shared_ptr<Camera> camera = active_camera_.lock();
        Vector3D loc = camera->GetActorLocate();
        eye = {loc.x, loc.y, loc.z};
        
        Vector3D rot = camera->GetActorRotate();
        glm::vec3 direction{glm::cos(rot.x)*glm::cos(rot.y), glm::sin(rot.x)*glm::cos(rot.y), glm::sin(rot.y)};
        center = eye + direction;
        
        up = {0, glm::sin(rot.z), glm::cos(rot.z)}; 
    }
    
    // static auto start_time = std::chrono::high_resolution_clock::now();
    // auto current_time = std::chrono::high_resolution_clock::now();
    // float time = std::chrono::duration_cast<std::chrono::duration<float>>(current_time - start_time).count();
    UniformBufferObject ubo;
    ubo.model = glm::rotate(glm::mat4(1.0f), /*time * glm::radians(45.0f)*/0.f, glm::vec3(0.0f, 0.0f, 1.0f));
    ubo.view = glm::lookAt(eye, center, up);
    ubo.proj = glm::perspective(glm::radians(45.0f), static_cast<float>(swapchain_.get_extent().width) / static_cast<float>(swapchain_.get_extent().height), 0.1f, 256.0f);
    ubo.proj[1][1] *= -1;
    memcpy(uniform_buffers_mapped_[current_frame], &ubo, sizeof(ubo));
}

void UwURenderEngine::draw_frame()
{
    vkWaitForFences(device_.get_native(), 1, &in_flight_fences_[current_frame_], VK_TRUE, UINT64_MAX);
    vkResetFences(device_.get_native(), 1, &in_flight_fences_[current_frame_]);
    
    // Free video memory
    garbage_collector_->FreeHeap();
    
    uint32_t image_index;
    vkAcquireNextImageKHR(device_.get_native(), swapchain_.get_native(), UINT64_MAX, image_available_semaphores_[current_frame_], VK_NULL_HANDLE, &image_index);
    
    vkResetCommandBuffer(command_buffers_[current_frame_], 0);
    
    update_uniform_buffer(current_frame_);
    record_command_buffer(command_buffers_[current_frame_], image_index);
    
    VkPipelineStageFlags wait_stage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    VkSubmitInfo submit_info{};
    submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submit_info.commandBufferCount = 1;
    submit_info.pCommandBuffers = &command_buffers_[current_frame_];
    submit_info.waitSemaphoreCount = 1;
    submit_info.pWaitSemaphores = &image_available_semaphores_[current_frame_];
    submit_info.pWaitDstStageMask = &wait_stage;
    submit_info.signalSemaphoreCount = 1;
    submit_info.pSignalSemaphores = &render_finished_semaphores_[current_frame_];
    
    VK_CHECK(vkQueueSubmit(device_.get_graphics_queue(), 1, &submit_info, in_flight_fences_[current_frame_]))
    
    VkSwapchainKHR swapchain = swapchain_.get_native();
    VkPresentInfoKHR present_info{};
    present_info.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
    present_info.waitSemaphoreCount = 1;
    present_info.pWaitSemaphores = &render_finished_semaphores_[current_frame_];
    present_info.swapchainCount = 1;
    present_info.pSwapchains = &swapchain;
    present_info.pImageIndices = &image_index;
    present_info.pResults = nullptr;
    
    vkQueuePresentKHR(device_.get_present_queue(), &present_info);
    
    current_frame_ = (current_frame_ + 1) % static_cast<unsigned int>(swapchain_.get_images().size());
}

UwURenderEngine::UwURenderEngine(CoreInstance& core):
RenderEngineBase(core),
core_(core),
instance_(core),
surface_(instance_, get_window()),
physical_device_(instance_, surface_, deviceExtensions),
device_(physical_device_, deviceExtensions),
swapchain_(surface_, physical_device_, device_, get_window()),
render_pass_(device_, swapchain_),
layout_binding_(device_),
pipeline_(device_, swapchain_, render_pass_, layout_binding_)
{
    create_uniform_buffers();
    create_descriptor_pool();
    create_descriptor_sets();
    create_framebuffers();
    create_command_pool();
    allocate_vertex_buffer();
    allocate_command_buffers();
    allocate_index_buffer();
    create_sync_objects();
    
    garbage_collector_.reset(new GPU_GarbageCollector(device_.get_native()));
    model_manager_.reset(new ModelManager(physical_device_.get_native(), device_.get_native()));
}

UwURenderEngine::~UwURenderEngine()
{
    vkDeviceWaitIdle(device_.get_native());
    
    if (index_buffer_memory_ != VK_NULL_HANDLE)
        vkFreeMemory(device_.get_native(), index_buffer_memory_, nullptr);
    
    if (index_buffer_ != VK_NULL_HANDLE)
        vkDestroyBuffer(device_.get_native(), index_buffer_, nullptr);
    
    if (vertex_buffer_memory_ != VK_NULL_HANDLE)
        vkFreeMemory(device_.get_native(), vertex_buffer_memory_, nullptr);

    if (vertex_buffer_ != VK_NULL_HANDLE)
        vkDestroyBuffer(device_.get_native(), vertex_buffer_, nullptr);
    
    for (auto& i : image_available_semaphores_)
        vkDestroySemaphore(device_.get_native(), i, nullptr);

    for (auto& i : render_finished_semaphores_)
        vkDestroySemaphore(device_.get_native(), i, nullptr);

    for (auto& i : in_flight_fences_)
        vkDestroyFence(device_.get_native(), i, nullptr);
    
    if (command_buffers_.empty() == false)
        vkFreeCommandBuffers(device_.get_native(), command_pool_, static_cast<uint32_t>(command_buffers_.size()), command_buffers_.data());

    if (command_pool_ != VK_NULL_HANDLE)
        vkDestroyCommandPool(device_.get_native(), command_pool_, nullptr);
        
    for (auto& i : framebuffers_)
        vkDestroyFramebuffer(device_.get_native(), i, nullptr);
    
    for (size_t i = 0; i < uniform_buffers_.size(); ++i)
    {
        vkUnmapMemory(device_.get_native(), uniform_buffers_memory_[i]);
        vkFreeMemory(device_.get_native(), uniform_buffers_memory_[i], nullptr);
        vkDestroyBuffer(device_.get_native(), uniform_buffers_[i], nullptr);
    }
    uniform_buffers_.clear();
    uniform_buffers_memory_.clear();
    uniform_buffers_mapped_.clear();
    
    vkDestroyDescriptorPool(device_.get_native(), descriptor_pool_, nullptr);
}

void UwURenderEngine::start()
{
    std::shared_ptr<World> world = core_.GetGameInstance()->GetWorld();
    std::vector<std::weak_ptr<Mesh>> meshes = world->GetActorsByClass<Mesh>();
    for (auto& i : meshes)
    {
        std::shared_ptr<Mesh> mesh = i.lock();
        mesh->load_model(mesh->GetModelName());
    }
        
    while (!glfwWindowShouldClose(get_window()))
    {
        glfwPollEvents();
        draw_frame();
    }
}

void UwURenderEngine::stop_render() const
{
    if (!glfwWindowShouldClose(get_window()))
        glfwSetWindowShouldClose(get_window(), true);
}

void UwURenderEngine::SetActiveCamera(const std::weak_ptr<Camera>& camera)
{
    std::scoped_lock lock(m_active_camera_);
    active_camera_ = camera;
}