/*
--------------------------------------------------------------------------------
                   ___
                  /  /
                 /  /
        ________/  /__________ _______
       /  _____/  /_______   //  ____/
      /  /	  /  /_______/  //  /
     /  /____/  //  _______//  /
    /__________//_____________/

    Ray Tracer Module for Screen-Space Reflections and GI
    
    Based on Belmu's NobleRT raytracer
    https://github.com/BelmuTM/NobleRT
    
    Upgraded from Shadax's MinecraftSSR
    https://github.com/Shadax-stack/MinecraftSSR

--------------------------------------------------------------------------------
*/

//============================================================================//
// RAYTRACER CONFIGURATION
//============================================================================//

#ifndef RAYTRACER_STEPS
    #define RAYTRACER_STEPS 64      // [32 64 128 256] Number of raytracing steps
#endif

#ifndef RAYTRACER_BISTEPS
    #define RAYTRACER_BISTEPS 4     // [0 2 4 8] Binary refinement steps
#endif

const uint rayTraceSteps = uint(RAYTRACER_STEPS);
const uint rayTraceBiSteps = uint(RAYTRACER_BISTEPS);
const float rayTracerStepsInv = 1.0 / float(RAYTRACER_STEPS);

//============================================================================//
// HELPER FUNCTIONS
//============================================================================//

// Fast normalize using inversesqrt
vec3 fastNormalize(in vec3 x) {
    return x * inversesqrt(dotSelf(x));
}

vec2 fastNormalize(in vec2 x) {
    return x * inversesqrt(dotSelf(x));
}

//============================================================================//
// MAIN RAYTRACER FUNCTION
//============================================================================//

/**
 * Performs screen-space ray tracing
 * 
 * @param screenPos    Current fragment screen position (0-1 range with depth)
 * @param viewPos      View-space position of current fragment
 * @param rayDir       View-space ray direction (normalized)
 * @param dither       Random dithering value [0-1] for temporal stability
 * 
 * @return vec3        xy = hit screen coordinates, z = hit success (1.0 = hit, 0.0 = miss)
 */
vec3 rayTraceScene(in vec3 screenPos, in vec3 viewPos, in vec3 rayDir, in float dither) {
    // Fix for the blob when player is near a surface
    // From Bálint#1673
    if (rayDir.z > -viewPos.z) return vec3(0.0);

    // Get screenspace rayDir
    // Convert view-space ray to screen-space
    vec3 screenPosRayDir = fastNormalize(ViewToScreenSpaceRaw(viewPos + rayDir) - screenPos) * rayTracerStepsInv;

    // Apply dithering to starting position for temporal stability
    vec3 startPos = screenPos + screenPosRayDir * dither;

    // Main raytracing loop
    for (uint i = 0u; i < rayTraceSteps; i++) {
        // Advance ray position
        startPos += screenPosRayDir;

        // If current pos is out of bounds, exit immediately
        if (startPos.x < 0.0 || startPos.y < 0.0 || startPos.x > 1.0 || startPos.y > 1.0) 
            return vec3(0.0);

        // Get current texture depth
        float currDepth = textureLod(depthtex0, startPos.xy, 0).x;

        // If hand return immediately (hand depth < 0.56 in Minecraft)
        if (currDepth <= 0.56) 
            return vec3(0.0);

        // Check intersection
        bool intersection = currDepth < startPos.z;

        // If intersection found
        if (intersection) {
            // Integrated binary refinement for precise hit location
            #if RAYTRACER_BISTEPS != 0
                for (uint j = 0u; j < rayTraceBiSteps; j++) {
                    // If sky return immediately
                    if (textureLod(depthtex0, startPos.xy, 0).x >= 1.0) 
                        return vec3(0.0);

                    // Continue refinement
                    screenPosRayDir *= 0.5;
                    startPos += intersection ? -screenPosRayDir : screenPosRayDir;

                    // Get current texture depth
                    currDepth = textureLod(depthtex0, startPos.xy, 0).x;
                    
                    // Check intersection
                    intersection = currDepth < startPos.z;
                }
            #else
                // If sky return immediately (no binary refinement)
                if (textureLod(depthtex0, startPos.xy, 0).x >= 1.0) 
                    return vec3(0.0);
            #endif

            // Return final results (xy = screen coords, z = success flag)
            return vec3(startPos.xy, 1.0);
        }
    }

    // No intersection found
    return vec3(0.0);
}

//============================================================================//
// ADDITIONAL RAYTRACING UTILITIES
//============================================================================//

/**
 * Screen-space ray tracer with Distant Horizons support
 * @param screenPos    Current fragment screen position
 * @param viewPos      View-space position
 * @param rayDir       View-space ray direction
 * @param dither       Dithering value
 * @param isDH         Whether to use Distant Horizons depth texture
 * @return vec3        Hit result (xy = coords, z = success)
 */
#ifdef DISTANT_HORIZONS
vec3 rayTraceSceneDH(in vec3 screenPos, in vec3 viewPos, in vec3 rayDir, in float dither) {
    if (rayDir.z > -viewPos.z) return vec3(0.0);

    vec3 screenPosRayDir = fastNormalize(ViewToScreenSpaceDH(viewPos + rayDir) - screenPos) * rayTracerStepsInv;
    vec3 startPos = screenPos + screenPosRayDir * dither;

    for (uint i = 0u; i < rayTraceSteps; i++) {
        startPos += screenPosRayDir;

        if (startPos.x < 0.0 || startPos.y < 0.0 || startPos.x > 1.0 || startPos.y > 1.0) 
            return vec3(0.0);

        // Use DH depth texture
        float currDepth = textureLod(dhDepthTex0, startPos.xy, 0).x;

        if (currDepth >= 1.0) 
            return vec3(0.0);

        bool intersection = currDepth < startPos.z;

        if (intersection) {
            #if RAYTRACER_BISTEPS != 0
                for (uint j = 0u; j < rayTraceBiSteps; j++) {
                    if (textureLod(dhDepthTex0, startPos.xy, 0).x >= 1.0) 
                        return vec3(0.0);

                    screenPosRayDir *= 0.5;
                    startPos += intersection ? -screenPosRayDir : screenPosRayDir;

                    currDepth = textureLod(dhDepthTex0, startPos.xy, 0).x;
                    intersection = currDepth < startPos.z;
                }
            #else
                if (textureLod(dhDepthTex0, startPos.xy, 0).x >= 1.0) 
                    return vec3(0.0);
            #endif

            return vec3(startPos.xy, 1.0);
        }
    }

    return vec3(0.0);
}
#endif

//============================================================================//
// ROUGH REFLECTIONS HELPER
//============================================================================//

/**
 * Generate cosine-weighted hemisphere direction for rough reflections
 * @param vector       Base vector (normal)
 * @param noiseVector  Random unit vector for perturbation
 * @return vec3        Perturbed direction
 */
vec3 generateCosineVector(in vec3 vector, in vec3 noiseVector) {
    vec3 vectorDir = fastNormalize(vector + noiseVector);
    return dot(vectorDir, vector) < 0.0 ? -vectorDir : vectorDir;
}

/**
 * Generate unit vector from 2D hash for random sampling
 * @param hash  2D random values [0-1]
 * @return vec3 Unit vector on sphere
 */
vec3 generateUnitVector(in vec2 hash) {
    hash.x *= TAU; 
    hash.y = hash.y * 2.0 - 1.0;
    return vec3(vec2(cos(hash.x), sin(hash.x)) * sqrt(1.0 - hash.y * hash.y), hash.y);
}
