#include <PipelineStage/DrawStage/CombineStage/CombineStage.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Vk/Framebuffer/Framebuffer.hpp>
#include <Vk/ShaderConst.hpp>
#include <Cubemap/Cubemap.hpp>
#include <Settings/Settings.hpp>

VkExtent2D CombineStage::getFramebufferExtent()
{
    return {
        (uint32_t) GraphicsContext::instance->getWindowWidth(),
        (uint32_t) GraphicsContext::instance->getWindowHeight(),
    };
}

CombineStage::CombineStage(
    FrameDrawer &frameDrawer,
    std::shared_ptr<Vk::DescriptorPool> combineDescriptorPool
) :
    DrawStage(
        "Combine",
        frameDrawer,
        *frameDrawer.worldDisplay.combineRenderPass,
        []() -> Vk::Framebuffer& { return GraphicsContext::instance->framebuffer(); },
        *frameDrawer.worldDisplay.combinePipeline,
        {}
    )
{
    combineDataBuffer = std::make_shared<Vk::Buffer>(
        GraphicsContext::instance->getDevice(),
        sizeof(CombineData),
        VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
    );

    skyboxDescSet = std::make_shared<Vk::DescriptorSet>(
        GraphicsContext::instance->getDevice(),
        combineDescriptorPool,
        frameDrawer.worldDisplay.skyboxDescriptorSetLayout
    );

    skyboxDescSet->setUniformBuffer(GLSL_SKYBOX_DATA_INDEX, *combineDataBuffer, 0, sizeof(CombineData));
    skyboxDescSet->setImageSampler(GLSL_SKYBOX_STAR_NOISE_INDEX, Cubemap::starNoise.imageView(), Vk::Sampler::cubemap());
    skyboxDescSet->setImageSampler(GLSL_SKYBOX_SUNMAP_INDEX, Cubemap::sunmap.imageView(), Vk::Sampler::cubemap());
}

void CombineStage::updateData()
{
    CombineData *combineData = (CombineData*) combineDataBuffer->hostdata();
    combineData->proj = frameDrawer.projMatrix;
    combineData->view = frameDrawer.worldDisplay.getPlayer()->getCameraTransform();
    combineData->model = frameDrawer.worldDisplay.getPlayer()->getSkyboxTransform();
    combineData->color = frameDrawer.worldDisplay.getWorld()->getSkyColor();
    combineData->screenSize = glm::vec2(GraphicsContext::instance->getWindowWidth(), GraphicsContext::instance->getWindowHeight());
    combineData->playerZ = frameDrawer.worldDisplay.getPlayer()->getEyePos().z;
    combineData->farZ = (float) (settings.horizontalViewDistance * 16);
    combineData->verticalViewDistance = (float) (settings.verticalViewDistance * 16);
    combineData->debugMode = frameDrawer.worldDisplay.gpuDebugMode;
    combineData->fxaaEnabled = settings.fxaa;
    combineData->godRaySamples = settings.godRaySamples;

    if (frameDrawer.worldDisplay.inThirdPersonView())
    {
        combineData->view = glm::mat4(
            -1.0f, 0.0f, 0.0f, 0.0f,
            0.0f, 1.0f, 0.0f, 0.0f,
            0.0f, 0.0f, -1.0f, 0.0f,
            0.0f, 0.0f, 0.0f, 1.0f
        ) * combineData->view;
    }

    cmdbuf->draw(6, 1, 0, 0);
}

void CombineStage::onFrameBegins()
{
    cmdbuf->bindDescriptorSet(
        VK_PIPELINE_BIND_POINT_GRAPHICS,
        pipeline,
        0,
        *skyboxDescSet
    );

    cmdbuf->bindDescriptorSet(
        VK_PIPELINE_BIND_POINT_GRAPHICS,
        pipeline,
        1,
        *frameDrawer.worldDisplay.combineDescriptorSet
    );
}