#define GLM_FORCE_RADIANS
#include <algorithm>
#include <chrono>
#include <cmath>
#include <filesystem>
#include <iostream>
#include <iterator>
#include <limits>
#include <set>
#include <random>

#include "Config.hpp"
#include "Core/CameraModel.hpp"
#include "Core/CgContext.hpp"
#include "Core/CoreInitException.hpp"
#include "Core/Model/AssimpModel.hpp"
#include "Core/Model/Object.hpp"
#include "Core/Model/SceneDescription.hpp"
#include "Util/Logging.hpp"
#include "Util/NameList.hpp"
#include "Util/PredefinedMeshes.hpp"
#include "Util/Utils.hpp"
#include "VkUtil/VkUtil.hpp"
#include "glm/glm.hpp"
#include "glm/gtc/matrix_transform.hpp"
#include "imgui_impl_vulkan.h"
#include "imgui_impl_glfw.h"

using lib::log::Debug;
using lib::log::Error;
using lib::log::Info;
using lib::log::Warning;

bool lib::CgContext::TryCreateRasterizerPipeline()
{
    auto success = true;
    auto vertexModule = this->LoadShader(this->appInfo.shaderTable.postProcessingVertexShader);

    success = success && vertexModule;

    auto fragmentModule = this->LoadShader(this->appInfo.shaderTable.postProcessingFragmentShader);

    success = success && fragmentModule;

    if (success)
    {
        std::vector<vk::PipelineShaderStageCreateInfo> shaderStages{
            vk::PipelineShaderStageCreateInfo{
                vk::PipelineShaderStageCreateFlags{},
                vk::ShaderStageFlagBits::eVertex,
                vertexModule,
                "main",
            },
            vk::PipelineShaderStageCreateInfo{
                vk::PipelineShaderStageCreateFlags{},
                vk::ShaderStageFlagBits::eFragment,
                fragmentModule,
                "main",
            },
        };

        vk::PipelineVertexInputStateCreateInfo emptyInputCreateInfo;

        vk::PipelineInputAssemblyStateCreateInfo inputAssemblyCreateInfo;
        inputAssemblyCreateInfo.topology = vk::PrimitiveTopology::eTriangleList;
        inputAssemblyCreateInfo.primitiveRestartEnable = false;

        vk::Viewport viewport;
        viewport.x = 0.f;
        viewport.y = 0.f;
        viewport.width = static_cast<float>(this->swapChainExtent.width);
        viewport.height = static_cast<float>(this->swapChainExtent.height);
        viewport.minDepth = 0.f;
        viewport.maxDepth = 1.f;

        vk::Rect2D scissor;
        scissor.offset = {0, 0};
        scissor.extent = this->swapChainExtent;

        vk::PipelineViewportStateCreateInfo viewportStateCreateInfo;
        viewportStateCreateInfo.viewportCount = 1;
        viewportStateCreateInfo.pViewports = &viewport;
        viewportStateCreateInfo.scissorCount = 1;
        viewportStateCreateInfo.pScissors = &scissor;

        vk::PipelineRasterizationStateCreateInfo rasterizationStateCreateInfo;
        rasterizationStateCreateInfo.depthClampEnable = false;
        rasterizationStateCreateInfo.rasterizerDiscardEnable = false;
        rasterizationStateCreateInfo.polygonMode = vk::PolygonMode::eFill;
        rasterizationStateCreateInfo.lineWidth = 1.f;
        rasterizationStateCreateInfo.frontFace = vk::FrontFace::eCounterClockwise;
        rasterizationStateCreateInfo.cullMode = vk::CullModeFlagBits::eBack;
        rasterizationStateCreateInfo.depthBiasEnable = false;

        vk::PipelineMultisampleStateCreateInfo multiSampleStageCreateInfo;
        multiSampleStageCreateInfo.sampleShadingEnable = false;
        multiSampleStageCreateInfo.rasterizationSamples = vk::SampleCountFlagBits::e1;

        vk::PipelineColorBlendAttachmentState colorBlendAttachmentState;
        colorBlendAttachmentState.colorWriteMask = vk::ColorComponentFlagBits::eR |
                                                   vk::ColorComponentFlagBits::eG |
                                                   vk::ColorComponentFlagBits::eB |
                                                   vk::ColorComponentFlagBits::eA;
        colorBlendAttachmentState.blendEnable = false;

        vk::PipelineColorBlendStateCreateInfo colorBlendStateCreateInfo;
        colorBlendStateCreateInfo.logicOpEnable = false;
        colorBlendStateCreateInfo.logicOp = vk::LogicOp::eCopy;
        colorBlendStateCreateInfo.attachmentCount = 1;
        colorBlendStateCreateInfo.pAttachments = &colorBlendAttachmentState;

        vk::PipelineLayoutCreateInfo layoutCreateInfo;
        layoutCreateInfo.setLayoutCount = 1;
        layoutCreateInfo.pSetLayouts = &this->rasterizationStep.descriptorSetLayout;

        this->rasterizationStep.pipelineLayout = this->logicalDevice.createPipelineLayout(layoutCreateInfo);

        vk::GraphicsPipelineCreateInfo createInfo;
        createInfo.stageCount = static_cast<uint32_t>(shaderStages.size());
        createInfo.pStages = shaderStages.data();
        createInfo.pVertexInputState = &emptyInputCreateInfo;
        createInfo.pInputAssemblyState = &inputAssemblyCreateInfo;
        createInfo.pViewportState = &viewportStateCreateInfo;
        createInfo.pRasterizationState = &rasterizationStateCreateInfo;
        createInfo.pMultisampleState = &multiSampleStageCreateInfo;
        createInfo.pColorBlendState = &colorBlendStateCreateInfo;
        createInfo.renderPass = this->rasterizationStep.renderPass;
        createInfo.layout = this->rasterizationStep.pipelineLayout;

        const auto result = this->logicalDevice.createGraphicsPipeline(nullptr, createInfo);

        if (result.result == vk::Result::eSuccess)
        {
            this->rasterizationStep.pipeline = result.value;
        }
        else
        {
            Error("Failed to create Graphicspipeline!");

            success = false;
        }
    }

    if (vertexModule)
    {
        this->logicalDevice.destroyShaderModule(vertexModule);
    }

    if (fragmentModule)
    {
        this->logicalDevice.destroyShaderModule(fragmentModule);
    }

    lib::log::Info("Successfully create rasterization pipeline!");

    return success;
}

void lib::CgContext::CreateRasterizationDescriptorSets()
{
    const auto count = static_cast<uint32_t>(this->frames.size());

    std::vector<vk::DescriptorSetLayout> layouts(count, this->rasterizationStep.descriptorSetLayout);

    vk::DescriptorSetAllocateInfo allocateInfo;
    allocateInfo.descriptorPool = this->rasterizationStep.descriptorPool;
    allocateInfo.descriptorSetCount = count;
    allocateInfo.pSetLayouts = layouts.data();

    auto sets = this->logicalDevice.allocateDescriptorSets(allocateInfo);
    auto setsIt = std::begin(sets);

    for (auto &frame : this->frames)
    {
        frame.rasterizationDescriptorSet = *setsIt++;

        vk::DescriptorImageInfo rtImageInfo;
        rtImageInfo.imageLayout = vk::ImageLayout::eShaderReadOnlyOptimal;
        // rtImageInfo.imageView = frame.raytracingTexture->GetView();
        // rtImageInfo.sampler = frame.raytracingTexture->GetSampler();

        rtImageInfo.imageView = this->dummyTexture->GetView();
        rtImageInfo.sampler = this->dummyTexture->GetSampler();

        vk::WriteDescriptorSet rtImageWrite;
        rtImageWrite.descriptorCount = 1;
        rtImageWrite.descriptorType = vk::DescriptorType::eCombinedImageSampler;
        rtImageWrite.dstArrayElement = 0;
        rtImageWrite.dstBinding = 0;
        rtImageWrite.pImageInfo = &rtImageInfo;
        rtImageWrite.dstSet = frame.rasterizationDescriptorSet;

        this->logicalDevice.updateDescriptorSets(rtImageWrite, {});
    }

    lib::log::Info("Successfully created rasterization descriptor sets!");
}

void lib::CgContext::CreateRasterizationDescriptorPool()
{
    const auto count = static_cast<uint32_t>(this->frames.size());

    std::vector<vk::DescriptorPoolSize> poolSizes{
        {vk::DescriptorType::eCombinedImageSampler, count},
    };

    vk::DescriptorPoolCreateInfo createInfo;
    createInfo.poolSizeCount = static_cast<uint32_t>(poolSizes.size());
    createInfo.pPoolSizes = poolSizes.data();
    createInfo.maxSets = count;

    this->rasterizationStep.descriptorPool = this->logicalDevice.createDescriptorPool(createInfo);

    lib::log::Info("Successfully created rasterization descriptor pool!");
}

void lib::CgContext::CreateRasterizationDescriptorSetLayout()
{
    std::vector<vk::DescriptorSetLayoutBinding> layoutBindings;

    auto &item = layoutBindings.emplace_back();
    item.binding = 0;
    item.descriptorType = vk::DescriptorType::eCombinedImageSampler;
    item.descriptorCount = 1;
    item.stageFlags = vk::ShaderStageFlagBits::eFragment;

    vk::DescriptorSetLayoutCreateInfo createInfo;
    createInfo.bindingCount = static_cast<uint32_t>(layoutBindings.size());
    createInfo.pBindings = layoutBindings.data();

    this->rasterizationStep.descriptorSetLayout = this->logicalDevice.createDescriptorSetLayout(createInfo);

    lib::log::Info("Successfully created rasterization descriptor set layouts!");
}

void lib::CgContext::CreateRasterizationRenderPass()
{
    vk::AttachmentDescription colorAttachment;
    colorAttachment.initialLayout = vk::ImageLayout::eUndefined;
    colorAttachment.finalLayout = vk::ImageLayout::ePresentSrcKHR;
    colorAttachment.loadOp = vk::AttachmentLoadOp::eDontCare;
    colorAttachment.storeOp = vk::AttachmentStoreOp::eStore;
    colorAttachment.format = this->swapChainImageFormat;
    colorAttachment.stencilLoadOp = vk::AttachmentLoadOp::eDontCare;
    colorAttachment.stencilStoreOp = vk::AttachmentStoreOp::eDontCare;

    vk::AttachmentReference colorAttachmentRef;
    colorAttachmentRef.layout = vk::ImageLayout::eColorAttachmentOptimal;
    colorAttachmentRef.attachment = 0;

    vk::SubpassDescription subpassDescription;
    subpassDescription.pipelineBindPoint = vk::PipelineBindPoint::eGraphics;
    subpassDescription.colorAttachmentCount = 1;
    subpassDescription.pColorAttachments = &colorAttachmentRef;

    vk::SubpassDependency subpassDependency;
    subpassDependency.srcSubpass = VK_SUBPASS_EXTERNAL;
    subpassDependency.dstSubpass = 0;
    subpassDependency.srcStageMask = vk::PipelineStageFlagBits::eColorAttachmentOutput;
    subpassDependency.srcAccessMask = vk::AccessFlagBits{};
    subpassDependency.dstStageMask = vk::PipelineStageFlagBits::eColorAttachmentOutput;
    subpassDependency.dstAccessMask = vk::AccessFlagBits::eColorAttachmentWrite;

    vk::RenderPassCreateInfo createInfo;
    createInfo.subpassCount = 1;
    createInfo.pSubpasses = &subpassDescription;
    createInfo.dependencyCount = 1;
    createInfo.pDependencies = &subpassDependency;
    createInfo.attachmentCount = 1;
    createInfo.pAttachments = &colorAttachment;

    this->rasterizationStep.renderPass = this->logicalDevice.createRenderPass(createInfo);

    lib::log::Info("Successfully created rasterization render pass!");
}

void lib::CgContext::CreateRasterizationFrameBuffers()
{
    const auto count = static_cast<uint32_t>(frames.size());

    for (auto &frame : this->frames)
    {
        const lib::TextureCreateInfo textureCreateInfo{
            {
                this->physicalDeviceInfo.GetGraphicsQueueFamily(),
                this->physicalDeviceInfo.GetBufferStagingQueueFamily(),
                this->physicalDeviceInfo.GetComputeQueueFamily(),
                this->physicalDeviceInfo.GetPresentationQueueFamily(),
            },
            this->swapChainImageFormat,
            vk::ImageUsageFlagBits::eSampled |
                vk::ImageUsageFlagBits::eStorage,
            vk::ImageLayout::eGeneral};

        auto newTexture = new lib::Texture{this,
                                           this->swapChainExtent.width,
                                           this->swapChainExtent.height,
                                           textureCreateInfo};

        frame.raytracingTexture.reset(newTexture);

        vk::FramebufferCreateInfo createInfo;
        createInfo.renderPass = this->rasterizationStep.renderPass;
        createInfo.attachmentCount = 1;
        createInfo.pAttachments = &frame.swapChainImageView;
        createInfo.width = this->swapChainExtent.width;
        createInfo.height = this->swapChainExtent.height;
        createInfo.layers = 1;

        frame.swapChainFramebufferForRasterizationRenderPass = logicalDevice.createFramebuffer(createInfo);
    }

    lib::log::Info("Successfully created rasterization framebuffers!");
}