#include <BlockRegistry/BlockRegistry.hpp>
#include <FrameDrawer/FrameDrawer.hpp>
#include <Display/WidgetDisplay/WorldDisplay/WorldDisplay.hpp>
#include <PipelineStage/PipelineStage.hpp>
#include <Settings/Settings.hpp>
#include <Vk/CommandBuffer/CommandBuffer.hpp>
#include <Vk/CommandPool/CommandPool.hpp>
#include <Vk/ShaderConst.hpp>
#include <Cubemap/Cubemap.hpp>
#include <IndexedUnitCube/IndexedUnitCube.hpp>
#include <ChunkWindow/ChunkWindow.hpp>
#include <memory>

void FrameDrawer::drawWorld(GraphicsContext &context)
{
    FrameData *frameData = (FrameData*) dataBuffer->hostdata();
    frameData->ambientColor = worldDisplay.world->getAmbientColor();
    frameData->solarAngle = worldDisplay.world->getSolarAngle();
    frameData->playerZ = worldDisplay.player->pos.z;
    frameData->cloudCover = worldDisplay.world->weather.getCloudCover();
    frameData->ticks = SDL_GetTicks();

    const float fovY = (float) settings.fov * (3.1415f / 180.0f);
    const float zNear = 0.1f;
    const float zFar = (float) (settings.horizontalViewDistance * 16);

    frameData->farZ = zFar;
    frameData->fogStartZ = zFar - fogLen;

    projMatrix = glm::perspective(
        fovY,
        (float) worldDisplay.framebufferExtent.width / (float) worldDisplay.framebufferExtent.height,
        zNear,
        zFar + 23.0f + 16.0f
    );
    projMatrix[1][1] *= -1;
    frameData->proj = projMatrix;

    viewMatrix = worldDisplay.player->getCameraTransform();
    if (worldDisplay.thirdPerson)
    {
        frameData->antiVP = glm::inverse(projMatrix * 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
        ) * viewMatrix);
        viewMatrix = glm::translate(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
        ), glm::vec3(0.0f, 0.0f, worldDisplay.player->getThirdPersonViewDistance())) * viewMatrix;
    }
    else
    {
        frameData->antiVP = glm::inverse(projMatrix * viewMatrix);
    }

    frameData->view = viewMatrix;

    frameData->skybox = worldDisplay.getPlayer()->getSkyboxTransform();
    frameData->antiSkybox = glm::inverse(frameData->skybox);

    frameData->cloudTransform = worldDisplay.cloudTransform;

    glm::dvec3 eye = worldDisplay.player->getEyePos();
    if (worldDisplay.inThirdPersonView())
    {
        eye += worldDisplay.player->getForwardVector() * worldDisplay.player->getThirdPersonViewDistance();
    }

    const glm::mat4 antiViewMatrix = glm::inverse(viewMatrix);
    const glm::dvec3 forward = antiViewMatrix * glm::vec4(0.0f, 0.0f, -1.0f, 0.0f);
    const glm::dvec3 up = antiViewMatrix * glm::vec4(0.0f, 1.0f, 0.0f, 0.0f);
    const glm::dvec3 right = antiViewMatrix * glm::vec4(1.0f, 0.0f, 0.0f, 0.0f);

    const double aspect = (float) context.getWindowWidth() / (float) context.getWindowHeight();

    const double halfVSide = tanf(fovY * .5f);
    const double halfHSide = halfVSide * aspect;

    glm::uvec3 eyeBlockPos = glm::uvec3(glm::ivec3(glm::floor(eye)));
    glm::vec3 blockVoxelOffset = 8.0f * glm::vec3(glm::fract(eye));

    frameData->cameraVoxel = (eyeBlockPos << 3u) + glm::uvec3(glm::ivec3(blockVoxelOffset));
    frameData->cameraOffset = glm::fract(blockVoxelOffset);

    frameData->flags = 0;
    
    if (settings.ambientOcclusion)
    {
        frameData->flags |= GLSL_FRAMEDATA_AMBIENT_OCCLUSION;
    }

    if (settings.shadows)
    {
        frameData->flags |= GLSL_FRAMEDATA_SHADOWS;
    }

    if (settings.specularLighting)
    {
        frameData->flags |= GLSL_FRAMEDATA_CONEMARCHING;
    }

    if (worldDisplay.enableDebugSpecular)
    {
        frameData->flags |= GLSL_FRAMEDATA_DEBUG_SPECULAR;
    }

    if (worldDisplay.enableDebugRefractValue)
    {
        frameData->flags |= GLSL_FRAMEDATA_DEBUG_REFRACT;
    }

    frameData->debugSpecular = (float) worldDisplay.debugSpecular / 255.0f;
    frameData->debugRefractValue = (float) worldDisplay.debugRefractValue / 255.0f;

    frameData->screenSize = glm::uvec2(worldDisplay.framebufferExtent.width, worldDisplay.framebufferExtent.height);

    frameData->screenTiles = (frameData->screenSize +
        glm::uvec2(GLSL_LIGHT_TILE_SIZE - 1)) /
        glm::uvec2(GLSL_LIGHT_TILE_SIZE);
    
    Frustum frustum;
    frustum.farFace = Plane::fromPointAndNormal(eye + (double) zFar * forward, -forward);
    frustum.leftFace = Plane::fromPointAndNormal(eye, glm::cross(forward - right * halfHSide, up));
    frustum.rightFace = Plane::fromPointAndNormal(eye, glm::cross(up, forward + right * halfHSide));
    frustum.topFace = Plane::fromPointAndNormal(eye, glm::cross(right, forward - up * halfVSide));
    frustum.bottomFace = Plane::fromPointAndNormal(eye, glm::cross(forward + up * halfVSide, right));

    worldDisplay.world->updateBeltAnimations(worldDisplay, frustum);
    worldDisplay.world->window->draw(worldDisplay);

    worldDisplay.world->drawBelts(worldDisplay, frustum);
    drawEntities(context);
    combineStage->updateData();
}

void FrameDrawer::drawEntities(GraphicsContext &context)
{
    // Draw the entities.
    worldDisplay.world->forEachEntity([this, &context](std::shared_ptr<Entity> entity)
    {
        entity->draw(context, worldDisplay);
    });
}

FrameDrawer::FrameDrawer(
    WorldDisplay &worldDisplay,
    std::shared_ptr<Vk::DescriptorPool> frameDataDescriptorPool,
    std::shared_ptr<Vk::DescriptorPool> skyboxDescriptorPool
) :
    worldDisplay(worldDisplay)
{
    std::shared_ptr<Vk::Device> device = GraphicsContext::instance->getDevice();

    dataBuffer = std::make_shared<Vk::Buffer>(
        device,
        sizeof(FrameData),
        VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
        VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT
    );

    dataDescSet = std::make_shared<Vk::DescriptorSet>(
        device,
        frameDataDescriptorPool,
        worldDisplay.worldFrameDescriptorSetLayout
    );

    dataDescSet->setUniformBuffer(0, *dataBuffer, 0, sizeof(FrameData));

    cmdpool = std::make_shared<Vk::CommandPool>(
        device,
        VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT | VK_COMMAND_POOL_CREATE_TRANSIENT_BIT,
        device->getQueueFamilies().graphicsFamily.value()
    );

    terrainStage = std::make_unique<TerrainStage>(*this, dataDescSet);
    stages.push_back(terrainStage.get());

    modelStage = std::make_unique<ModelStage>(*this, dataDescSet);
    stages.push_back(modelStage.get());

    particlesStage = std::make_unique<ParticlesStage>(*this, dataDescSet);
    stages.push_back(particlesStage.get());

    lightingStage = std::make_unique<LightingStage>(*this, dataDescSet, frameDataDescriptorPool);
    stages.push_back(lightingStage.get());

    combineStage = std::make_unique<CombineStage>(*this, skyboxDescriptorPool);
    stages.push_back(combineStage.get());
}

void FrameDrawer::draw(GraphicsContext &context)
{
    for (PipelineStage *stage : stages)
    {
        stage->initFrame();
    }

    drawWorld(context);

    for (PipelineStage *stage : stages)
    {
        stage->encodeInto(context.cmdbuf());
    }
}

std::shared_ptr<Vk::CommandBuffer> FrameDrawer::createStageCommandBuffer()
{
    return std::make_shared<Vk::CommandBuffer>(
        GraphicsContext::instance->getDevice(),
        cmdpool,
        VK_COMMAND_BUFFER_LEVEL_SECONDARY
    );
}

void FrameDrawer::setModelMatrix(const glm::mat4 &modelMatrix)
{
    modelStage->setModelMatrix(modelMatrix);
}

void FrameDrawer::writeStageDebugLines(std::ostream &os)
{
    for (PipelineStage *stage : stages)
    {
        os << "  " << stage->label << ": " << stage->timeTracker.getAverage() << "ms" << std::endl;
    }
}