// TODO: Integrate with Blender UI
// TODO: Option to disable cone tracing

#version 450

#include "compiled.inc"
#include "std/math.glsl"
#include "std/gbuffer.glsl"

// Environment map
uniform float envmapStrength;
#ifdef _Irr
	uniform vec4 shirr[7];
	uniform float ambientIntensity;
    #include "std/shirr.glsl"
#endif
#include "std/environment_sample.glsl"

uniform sampler2D tex;                // Environment map
//uniform sampler2D depthtex;           // Full Depth buffer
uniform sampler2D gbufferD;           // Cheap Depth buffer
uniform sampler2D gbuffer0;           // Normal, roughness
uniform sampler2D gbuffer1;           // Base color, spec
uniform mat4 P;                       // Projection matrix
uniform mat3 V3;                      // View matrix
uniform vec2 cameraProj;              // Camera projection params
uniform vec2 invScreenSize;           // (1.0/width, 1.0/height)

const float ssrPrecision = 0.0;     // 0.0 - 100.0 (user slider control)
//const float rayThickness = 0.1;      // TODO: Adds some thickness to prevent gaps
uniform int ssrConetraceMode = 2;          // 0 = no weighting, 1 = light, 2 = strong
const int ssrConetraceTaps = 18;     // Number of taps (higher =  more precision)

in vec3 viewRay;
in vec2 texCoord;
out vec4 fragColor;

vec3 hitCoord;
float depth;

const int baseBinarySearchSteps = 10; // Parameterize?
const int baseMaxSteps = int(ceil(1.0 / ssrRayStep) * ssrSearchDist);

int dynamicBinarySearchSteps() {
    return int(mix(7.0, 20.0, clamp(ssrPrecision / 100.0, 0.0, 1.0)));
}

int dynamicMaxSteps() {
    return int(mix(float(baseMaxSteps), 300.0, clamp(ssrPrecision / 100.0, 0.0, 1.0)));
}

vec2 getProjectedCoord(const vec3 hit) {
    vec4 clip = P * vec4(hit, 1.0);
    vec2 uv   = clip.xy / clip.w;
    uv = uv * 0.5 + 0.5;
#ifdef _InvY
    uv.y = 1.0 - uv.y;
#endif
    uv = clamp(uv, 0.0, 1.0);
    uv += invScreenSize * 0.5; // half-pixel offset
    return uv;
}

float getDeltaDepth(const vec3 hit) {
    vec2 screenUV = getProjectedCoord(hit);
    float raw = textureLod(gbufferD, screenUV, 0.0).r;
    float sampledDepth = raw * 2.0 - 1.0;
    float linearSampledDepth = linearize(sampledDepth, cameraProj);

    return linearSampledDepth - hit.z;
}

vec4 binarySearch(vec3 dir) {
    float ddepth;
    for (int i = 0; i < dynamicBinarySearchSteps(); i++) {
        dir *= 0.5;
        hitCoord -= dir;
        ddepth = getDeltaDepth(hitCoord);
        if (ddepth < 0.0) hitCoord += dir;
    }

    vec2 projectedCoord = getProjectedCoord(hitCoord);
    float pixelSize = length(fwidth(projectedCoord)) * 0.5;
    float epsilon = max(pixelSize * 10.0, 0.01);

    if (abs(ddepth) > epsilon) return vec4(0.0);

    hitCoord.xy = clamp(projectedCoord, 0.0, 1.0);

    return vec4(projectedCoord, 0.0, 1.0);
}

vec4 rayCast(vec3 dir, float roughness) {
    vec3 stepDir = normalize(dir);

    // Apply small jitter if high precision
    if (ssrPrecision > 80.0 && roughness > 0.05) {
		stepDir = normalize(stepDir + vec3(rand2(texCoord), 0.0) * 0.01);
	}

    float distance = length(hitCoord - viewRay);
    float stepDivisor = mix(100.0, 300.0, clamp(ssrPrecision / 100.0, 0.0, 1.0));
    float stepSize = max(0.01, distance / stepDivisor);
    float maxStepSize = 0.1;

    for (int i = 0; i < dynamicMaxSteps(); i++) {
        hitCoord += stepDir * stepSize;

        vec2 projCoord = getProjectedCoord(hitCoord);
        float sampledDepth = textureLod(gbufferD, projCoord, 0.0).r;
        float depthTolerance = fwidth(sampledDepth) * 2.0; // Dynamic tolerance

        if (getDeltaDepth(hitCoord) > depthTolerance) {
            return binarySearch(stepDir);
        }

        stepSize = min(stepSize * 1.1, maxStepSize);
    }

    return vec4(0.0);
}

vec3 coneTraceApprox(vec3 reflDir, float roughness) {
    vec3 result = vec3(0.0);
    float totalWeight = 0.0;

    float randAngle = hash12(texCoord) * 6.2831853;
    mat2 rot = mat2(cos(randAngle), -sin(randAngle), sin(randAngle), cos(randAngle));

    float coneSpread = roughness * 0.5; // widen based on roughness

    for (int i = 0; i < ssrConetraceTaps; ++i) {
        float angle = float(i) / float(ssrConetraceTaps) * 6.2831853;
        vec2 offset = rot * vec2(cos(angle), sin(angle)) * coneSpread;
        vec3 sampleDir = normalize(reflDir + vec3(offset, 0.0));

        float weight = 1.0;
        if (ssrConetraceMode == 1) {
            weight = pow(max(dot(reflDir, sampleDir), 0.0), 2.0); // Light cosine lobe
        }
        else if (ssrConetraceMode == 2) {
            weight = pow(max(dot(reflDir, sampleDir), 0.0), 8.0); // Strong GGX lobe
        }

        // Approximate sampling by using tex at reflection direction
        vec2 envUV = envMapEquirect(sampleDir);
        vec3 sampleColor = textureLod(tex, envUV, 0.0).rgb;

        result += sampleColor * weight;
        totalWeight += weight;
    }

    return result / max(totalWeight, 0.0001);
}

vec3 tangentSpaceGGX(vec3 N, float roughness) {
    float a = roughness * roughness;
    float phi = rand(texCoord) * 6.2831853;
    float cosTheta = sqrt((1.0 - rand(texCoord)) / (1.0 + (a * a - 1.0) * rand(texCoord)));
    float sinTheta = sqrt(1.0 - cosTheta * cosTheta);

    vec3 T = normalize(cross(N, vec3(0.0, 1.0, 0.0)));
    if (length(T) < 0.01) T = normalize(cross(N, vec3(1.0, 0.0, 0.0)));
    vec3 B = cross(N, T);

    return normalize(T * cos(phi) * sinTheta + B * sin(phi) * sinTheta + N * cosTheta);
}

void main() {

    vec4 g0 = textureLod(gbuffer0, texCoord, 0.0);
    float roughness = unpackFloat(g0.b).y;
    if (roughness == 1.0) { fragColor.rgb = vec3(0.0); return; }

    float spec = fract(textureLod(gbuffer1, texCoord, 0.0).a);
    if (spec == 0.0) { fragColor.rgb = vec3(0.0); return; }

    // sample raw depth, bail if empty
    float dRaw = textureLod(gbufferD, texCoord, 0.0).r;
    if (dRaw == 0.0) { fragColor.rgb = vec3(0.0); return; }
    // convert to NDC z before reconstructing
    float d = dRaw * 2.0 - 1.0;

    vec3 n = decode_oct(g0.rg);
    vec3 viewNormal = V3 * n;
    viewNormal = normalize(mix(viewNormal, normalize(-viewRay), 0.05)); // slightly bias the normal toward the view direction at glancing angles
    vec3 viewPos = getPosView(viewRay, d, cameraProj);
    vec3 viewDir = normalize(-viewPos);
    vec3 idealReflection = reflect(normalize(viewPos), viewNormal);
    hitCoord = viewPos;

    // Apply GGX importance sampling in tangent space
    vec3 jitteredDir = tangentSpaceGGX(viewNormal, roughness);

    // Blend based on roughness (0 = perfect mirror, 1 = fully scattered)
    vec3 dir = normalize(mix(idealReflection, jitteredDir, roughness));

    vec4 coords = rayCast(dir, roughness);

    vec2 deltaCoords = abs(vec2(0.5, 0.5) - coords.xy);
    float screenEdgeFactor = clamp(1.0 - (deltaCoords.x + deltaCoords.y), 0.0, 1.0);

    float reflectivity = 1.0 - roughness;

    #ifdef _CPostprocess
		// Postprocess mode intensity calculation...
	#else
		float intensity = 0.0;

		if (coords.w > 0.0) {
		    // Ray hit in screen-space
		    intensity = pow(reflectivity, ssrFalloffExp) * screenEdgeFactor *
		                clamp(-idealReflection.z, 0.0, 1.0) *
		                clamp((ssrSearchDist - length(viewPos - hitCoord)) * (1.0 / ssrSearchDist), 0.0, 1.0);
		}
		else if (roughness < 0.7) {
		    // Ray miss, roughness low enough: do cone trace approximation
		    vec3 coneColor = coneTraceApprox(idealReflection, roughness);
		    vec3 envCol = sampleSpecularEnvironment(idealReflection, roughness);

		    coneColor = clamp(coneColor, 0.0, 1.0);
		    envCol = clamp(envCol, 0.0, 1.0);

		    vec3 finalColor = mix(envCol, mix(envCol, coneColor, intensity), coords.w);
		    fragColor.rgb = finalColor * 0.5; // match previous intensity scaling
		    return;
		}
		else {
		    // Ray miss, roughness too high: fallback to environment map
		    vec3 fallbackEnvColor = sampleSpecularEnvironment(reflect(viewDir, n), roughness);
		    fragColor.rgb = fallbackEnvColor;
		    return;
		}
	#endif

    vec3 reflCol = textureLod(tex, coords.xy, 0.0).rgb; // SSR reflection
	vec3 envCol = textureLod(tex, texCoord, 0.0).rgb;    // Background environment

    intensity = clamp(intensity, 0.0, 1.0);
	reflCol = clamp(reflCol, 0.0, 1.0);
	envCol = clamp(envCol, 0.0, 1.0);

	// Roughness-based fade (smoothstep)
	float roughFade = smoothstep(0.2, 0.7, roughness);
	float ssrVisibility = coords.w * (1.0 - roughFade);

	// Blend SSR
	vec3 finalColor = mix(envCol, reflCol, ssrVisibility);

	fragColor.rgb = finalColor;
}