#version 460

#extension GL_GOOGLE_include_directive : require
#extension GL_EXT_nonuniform_qualifier : enable

#include "bindless.glsl"

// The selection image is R32_UINT, read through the same bindless array as an unsigned texture
layout (set = 0, binding = 0) uniform utexture2D uintTextures[];

layout (location = 0) in vec2 inUV;
layout (location = 0) out vec4 outColor;

// Must match PostProcessPipeline::PushConstants and the Tonemapper enum
#define TONEMAPPER_NONE 0
#define TONEMAPPER_ACES 1
#define TONEMAPPER_PBR_NEUTRAL 2
#define NO_TEXTURE 0xFFFFFFFFu
#define SELECTED_FLAG 0x80000000u

layout (push_constant) uniform PostPushConstants {
    uint sceneColorTexture;
    uint selectionTexture;
    uint tonemapper;
    float exposure;
    ivec2 hoverPixel;
} pcs;

// Narkowicz 2015, a fit of the ACES filmic curve
vec3 tonemapAces(vec3 color)
{
    const float a = 2.51;
    const float b = 0.03;
    const float c = 2.43;
    const float d = 0.59;
    const float e = 0.14;
    return clamp((color * (a * color + b)) / (color * (c * color + d) + e), 0.0, 1.0);
}

// Khronos PBR Neutral, keeps base colors as they are up to about 0.8 and only rolls off the highlights
vec3 tonemapPbrNeutral(vec3 color)
{
    const float startCompression = 0.8 - 0.04;
    const float desaturation = 0.15;

    float x = min(color.r, min(color.g, color.b));
    float offset = x < 0.08 ? x - 6.25 * x * x : 0.04;
    color -= offset;

    float peak = max(color.r, max(color.g, color.b));
    if (peak < startCompression) {
        return color;
    }

    const float d = 1.0 - startCompression;
    float newPeak = 1.0 - d * d / (peak + d - startCompression);
    color *= newPeak / peak;

    float g = 1.0 - 1.0 / (desaturation * (peak - newPeak) + 1.0);
    return mix(color, vec3(newPeak), g);
}

uint objectAt(ivec2 pixel, ivec2 size)
{
    pixel = clamp(pixel, ivec2(0), size - 1);
    return texelFetch(nonuniformEXT(usampler2D(uintTextures[pcs.selectionTexture], samplers[NEAREST_SAMPLER_ID])), pixel, 0).r;
}

vec3 applyOutlines(vec3 color, ivec2 pixel)
{
    if (pcs.selectionTexture == NO_TEXTURE) {
        return color;
    }

    ivec2 size = textureSize(nonuniformEXT(usampler2D(uintTextures[pcs.selectionTexture], samplers[NEAREST_SAMPLER_ID])), 0);
    uint center = objectAt(pixel, size);
    uint hovered = pcs.hoverPixel.x >= 0 ? (objectAt(pcs.hoverPixel, size) & ~SELECTED_FLAG) : 0u;

    bool selectedEdge = false;
    bool hoveredEdge = false;
    for (int y = -2; y <= 2; ++y) {
        for (int x = -2; x <= 2; ++x) {
            if (abs(x) + abs(y) > 3) {
                continue;
            }
            uint neighbor = objectAt(pixel + ivec2(x, y), size);
            selectedEdge = selectedEdge || ((neighbor & SELECTED_FLAG) != 0u && (center & SELECTED_FLAG) == 0u);
            hoveredEdge = hoveredEdge || (hovered != 0u && (neighbor & ~SELECTED_FLAG) == hovered &&
                                          (center & ~SELECTED_FLAG) != hovered);
        }
    }

    if (selectedEdge) {
        return vec3(1.0, 0.55, 0.1);
    }
    if (hoveredEdge) {
        return mix(color, vec3(1.0), 0.65);
    }
    return color;
}

void main()
{
    ivec2 pixel = ivec2(gl_FragCoord.xy);
    vec3 color = texelFetch(nonuniformEXT(sampler2D(textures[pcs.sceneColorTexture], samplers[NEAREST_SAMPLER_ID])), pixel, 0).rgb;
    color *= pcs.exposure;

    if (pcs.tonemapper == TONEMAPPER_ACES) {
        color = tonemapAces(color);
    } else if (pcs.tonemapper == TONEMAPPER_PBR_NEUTRAL) {
        color = tonemapPbrNeutral(color);
    }

    outColor = vec4(applyOutlines(clamp(color, 0.0, 1.0), pixel), 1.0);
}