/* Aurora Fantasy - water.glsl
Water reflection and refraction related functions.
*/

#include "/lib/water_waves.glsl"
#include "/lib/water_palette.glsl"

vec3 fast_raymarch(
    vec3 direction,
    vec3 hit_coord,
    inout float infinite,
    out float hit_confidence,
    float dither
) {
    vec3 dir_increment;
    vec3 current_march = hit_coord;
    vec3 old_current_march;
    float screen_depth = 1.0;
    float depth_diff = 1.0;
    vec3 march_pos = camera_to_screen(hit_coord);
    float prev_screen_depth = march_pos.z;
    float hit_z = march_pos.z;
    bool search_flag = false;
    bool hidden_flag = false;
    bool out_flag = false;
    bool to_far = false;
    bool valid_hit = false;
    vec3 last_march_pos;
    
    int no_hidden_steps = 0;
    bool hiddens = false;

    hit_confidence = 0.0;

    // A fixed iteration budget is selected by the shader profile. Changing
    // the stride with instantaneous FPS made reflections pop between several
    // different ray paths and did not preserve their energy.
    for (int i = 0; i < RAYMARCH_STEPS; i++) {
        if (search_flag) {
            dir_increment *= 0.5;
            current_march += dir_increment * sign(depth_diff);
        } else {
            old_current_march = current_march;
            // Spread every quality tier over the same useful world-space
            // range. More samples now add precision instead of extending the
            // march into depths that cannot produce a visible reflection.
            float march_exponent =
                (float(i) + dither)
                * (10.0 / float(RAYMARCH_STEPS - 1));
            current_march =
                hit_coord + direction * (exp2(march_exponent) - 1.0);
            dir_increment = current_march - old_current_march;
        }

        last_march_pos = march_pos;
        march_pos = camera_to_screen(current_march);

        if ( // Is outside screen space
            march_pos.x < 0.0 ||
            march_pos.x > RENDER_SCALE ||
            march_pos.y < 0.0 ||
            march_pos.y > RENDER_SCALE ||
            march_pos.z < 0.0
        ) {
            out_flag = true;
            break;
        }

        if (march_pos.z > 0.9999) {
            to_far = true;
        }

        screen_depth = texture2D(depthtex1, march_pos.xy).x;
        depth_diff = screen_depth - march_pos.z;

        if (depth_diff < 0.0 && abs(screen_depth - prev_screen_depth) > abs(march_pos.z - last_march_pos.z)) {
            hidden_flag = true;
            hiddens = true;
        } else if (depth_diff > 0.0) {
            hidden_flag = false;
            if (!hiddens) {
                no_hidden_steps++;
            }
        }

        if (search_flag == false && depth_diff < 0.0 && hidden_flag == false) {
            search_flag = true;
            valid_hit = true;
        }

        prev_screen_depth = screen_depth;
    }

    infinite = float(!valid_hit || screen_depth > 0.9999 || out_flag || to_far);

    if (out_flag) {
        return clamp(march_pos, vec3(0.0), vec3(RENDER_SCALE, RENDER_SCALE, 1.0));
    } else if (to_far) {
        valid_hit = valid_hit && screen_depth < 0.9999
            && (no_hidden_steps < 3 || screen_depth > hit_z);
    }

    march_pos.xy = clamp(march_pos.xy, vec2(0.0), vec2(RENDER_SCALE));
    hit_confidence = float(valid_hit);
    infinite = 1.0 - hit_confidence;
    return march_pos;
}

#if SUN_REFLECTION > 0
    #if !defined NETHER && !defined THE_END
        float sun_reflection(vec3 fragpos, float smoothstep1) {
            vec3 nfragPos = normalize(fragpos);
            vec3 astro_pos = worldTime > 13200 ? moonPosition : sunPosition;
            vec3 nastroPos = normalize(astro_pos);

            float reflection = 0.0;

            #if  SUN_REFLECTION == 2
                vec3 forward = normalize(mat3(gbufferModelViewInverse) * nastroPos);
                vec3 nFragWorld = normalize(mat3(gbufferModelViewInverse) * nfragPos);

                vec3 worldRef = vec3(0.0, 1.0, 0.0);
                if (abs(dot(forward, worldRef)) > 0.99) worldRef = vec3(0.0, 0.0, 1.0);

                vec3 rightBase = normalize(cross(worldRef, forward));
                vec3 upBase = cross(forward, rightBase);

                float rawX = dot(nFragWorld, rightBase);
                float rawY = dot(nFragWorld, upBase);

                float angle = sunPathRotation * 0.01745329; 
                float s = sin(angle);
                float c = cos(angle);
                
                // Rotation
                float dotX = rawX * c - rawY * s;
                float dotY = rawX * s + rawY * c;
                
                float square_shape = max(abs(dotX), abs(dotY));
                float size = 1.08 - smoothstep1;
                float frontal_mask = step(0.0, dot(nFragWorld, forward));
                
                reflection = smoothstep(size + 0.01, size, square_shape) * frontal_mask;
            #else
                float astro_vector = max(dot(nfragPos, nastroPos), 0.0);
                reflection = smoothstep(smoothstep1, 1.0, astro_vector);
            #endif

            return clamp(
                (reflection - 0.165) *
                clamp(lmcoord.y, 0.0, 1.0) *
                (1.0 - rainStrength) * day_blend_float(10.0, 10.0, 5.0), 0.0, 100.0);
        }
    #endif
#endif

vec3 normal_waves(vec3 pos) {
    return auroraWaterNormalWaves(
        pos, persistentTimeSeconds, rainStrength, visible_sky);
}

vec3 refraction(vec3 fragpos, vec3 color, vec3 refraction) {
    vec2 pos = gl_FragCoord.xy * vec2(pixel_size_x, pixel_size_y);

    #if REFRACTION == 1
        pos = pos + refraction.xy * (0.075 / (1.0 + length(fragpos) * 0.4));
    #endif

    float water_absortion;
    if (isEyeInWater == 0) {
        float water_distance =
        2.0 * near * far / (far + near - (2.0 * gl_FragCoord.z - 1.0) * (far - near));

        float earth_distance = texture2D(depthtex1, pos.xy).r;
        earth_distance =
            2.0 * near * far / (far + near - (2.0 * earth_distance - 1.0) * (far - near));

        #if defined DISTANT_HORIZONS
            float earth_distance_dh = texture2D(dhDepthTex1, pos.xy).r;
            earth_distance_dh =
                2.0 * dhNearPlane * dhFarPlane / (dhFarPlane + dhNearPlane - (2.0 * earth_distance_dh - 1.0) * (dhFarPlane - dhNearPlane));
            earth_distance = min(earth_distance, earth_distance_dh);
        #endif

        water_absortion = earth_distance - water_distance;
        water_absortion *= water_absortion;
        water_absortion = (1.0 / -((water_absortion * WATER_ABSORPTION) + 1.33)) + 1.0;
    } else {
        water_absortion = 0.0;
    }

    return mix(texture2D(gaux1, pos.xy).rgb, color, water_absortion);
}

vec3 get_normals(vec3 bump, vec3 fragpos) {
    float NdotE = abs(dot(water_normal, normalize(fragpos)));

    bump *= vec3(NdotE) + vec3(0.05, 0.1, 1.0 - NdotE);

    mat3 tbn_matrix = mat3(
        tangent.x, binormal.x, water_normal.x,
        tangent.y, binormal.y, water_normal.y,
        tangent.z, binormal.z, water_normal.z
    );

    return normalize(bump * tbn_matrix);
}

vec4 reflection_calc(vec3 fragpos, vec3 normal, vec3 reflected, inout float infinite, float dither) {
    float hit_confidence = 1.0;
    #if SSR_TYPE == 0  // Flipped image
        #if defined DISTANT_HORIZONS
            vec3 pos = camera_to_screen(fragpos + reflected * 768.0);
        #else
            vec3 pos = camera_to_screen(fragpos + reflected * 76.0);
        #endif
    #else  // Raymarch
        vec3 pos = fast_raymarch(reflected, fragpos, infinite, hit_confidence, dither);
    #endif
    float pos_y_normalized = pos.y / RENDER_SCALE;

    float border =
        clamp((1.0 - (max(0.0, abs(pos_y_normalized - 0.5)) * 2.0)) * 50.0, 0.0, 1.0);

    border = clamp(border - pow(pos_y_normalized, 10.0), 0.0, 1.0);

    pos.x = abs(pos.x);

    if (pos.x > RENDER_SCALE) {
        pos.x = RENDER_SCALE - (pos.x - RENDER_SCALE);
    }
    
    return vec4(texture2D(gaux1, pos.xy).rgb, border * hit_confidence);
}

vec3 water_shader(
    vec3 fragpos,
    vec3 normal,
    vec3 color,
    vec3 sky_reflect,
    vec3 reflected,
    float fresnel,
    float visible_sky,
    float dither,
    vec3 light_color
) {
    vec4 reflection = vec4(0.0);
    float infinite = 1.0;

    #if REFLECTION == 1
        reflection =
            reflection_calc(fragpos, normal, reflected, infinite, dither);
    #endif

    reflection.rgb = mix(
        sky_reflect * visible_sky,
        reflection.rgb,
        reflection.a
    );

    #ifdef VANILLA_WATER
        fresnel *= 0.8;
    #endif

    #if SUN_REFLECTION > 0
        #ifndef NETHER
            #ifndef THE_END
                return mix(color, reflection.rgb, fresnel * REFLEX_INDEX) +
                    vec3(sun_reflection(reflect(normalize(fragpos), normal), 0.999)) * light_color * infinite * visible_sky * day_blend(vec3(1.0, 1.0, 0.35), vec3(1.0, 1.0, 0.35), vec3(1.0));          
            #else
                return mix(color, reflection.rgb, fresnel * REFLEX_INDEX);
            #endif
        #else
            return mix(color, reflection.rgb, fresnel * REFLEX_INDEX);
        #endif
    #else
        return mix(color, reflection.rgb, fresnel * REFLEX_INDEX);
    #endif
}

//  GLASS

vec4 cristal_reflection_calc(vec3 fragpos, vec3 normal, inout float infinite, float dither) {
    float hit_confidence = 1.0;
    #if SSR_TYPE == 0
        #if defined DISTANT_HORIZONS
            vec3 reflected_vector = reflect(normalize(fragpos), normal) * 768.0;
        #else
            vec3 reflected_vector = reflect(normalize(fragpos), normal) * 76.0;
        #endif
            vec3 pos = camera_to_screen(fragpos + reflected_vector);
    #else
        vec3 reflected_vector = reflect(normalize(fragpos), normal);
        vec3 pos = fast_raymarch(
            reflected_vector, fragpos, infinite, hit_confidence, dither);
    #endif

        vec2 normalized_pos = pos.xy / RENDER_SCALE;
        float border_x = max(-fourth_pow(abs(2.0 * normalized_pos.x - 1.0)) + 1.0, 0.0);
        float border_y = max(-fourth_pow(abs(2.0 * normalized_pos.y - 1.0)) + 1.0, 0.0);
        float border = min(border_x, border_y);
    
    // A missed SSR ray contributes no screen image; cristal_shader then uses
    // its physically stable sky reflection instead of a projected fake hit.
    return vec4(texture2D(gaux1, pos.xy).rgb, border * hit_confidence);
}

vec4 cristal_shader(
    vec3 fragpos,
    vec3 normal,
    vec4 color,
    vec3 sky_reflection,
    float fresnel,
    float visible_sky,
    float dither,
    vec3 light_color
) {
    vec4 reflection = vec4(0.0);
    float infinite = 0.0;

    #if REFLECTION == 1
        reflection = cristal_reflection_calc(fragpos, normal, infinite, dither);
    #endif

    sky_reflection = mix(color.rgb, sky_reflection, visible_sky * visible_sky);

    reflection.rgb = mix(
        sky_reflection,
        reflection.rgb,
        reflection.a
    );

    color.rgb = mix(color.rgb, sky_reflection, fresnel);
    color.rgb = mix(color.rgb, reflection.rgb, fresnel);

    color.a = mix(color.a, 1.0, fresnel);

    #if SUN_REFLECTION > 0
        #ifndef NETHER
            #ifndef THE_END
                return color + vec4(
                    mix(
                            vec3(sun_reflection(reflect(normalize(fragpos), normal), 0.99) * 0.1 * light_color * infinite * day_blend(vec3(1.0, 1.0, 0.15), vec3(1.0, 1.0, 0.15), vec3(1.0))),
                        vec3(0.0),
                        reflection.a
                    ),
                    0.0
                );
            #else
            return color;
        #endif
        #else
            return color;
        #endif
    #else
        return color;
    #endif
}
