#version 450 core
#include "common.glsl"

in vec3 vWorld;

layout(binding = 0) uniform sampler2D uReflection;
layout(binding = 1) uniform sampler2D uSceneColor;
layout(binding = 2) uniform sampler2D uSceneDepth;
layout(binding = 3) uniform sampler2D uShoreDist;

uniform vec4 uShoreRect;  // minX, minZ, 1/sizeX, 1/sizeZ
uniform vec3 uShallow;
uniform vec3 uDeep;
uniform vec3 uAbsorb;
uniform float uClarity;
uniform float uWaveScale;
uniform float uWaveStrength;
uniform float uWaveSpeed;
uniform float uRefraction;
uniform float uReflDistort;
uniform float uSpecular;
uniform float uGlitter;
uniform float uFoamWidth;
uniform float uFoamIntensity;
uniform int uHasReflection;

layout(location = 0) out vec4 fragColor;

// Value noise with analytic derivatives (iq).
vec3 noised(vec2 x) {
    vec2 i = floor(x), f = fract(x);
    vec2 u = f * f * f * (f * (f * 6.0 - 15.0) + 10.0);
    vec2 du = 30.0 * f * f * (f * (f - 2.0) + 1.0);
    float a = hash12(i), b = hash12(i + vec2(1, 0)), c = hash12(i + vec2(0, 1)), d = hash12(i + vec2(1, 1));
    float k1 = b - a, k2 = c - a, k4 = a - b - c + d;
    return vec3(a + k1 * u.x + k2 * u.y + k4 * u.x * u.y, du * vec2(k1 + k4 * u.y, k2 + k4 * u.x));
}

// Twinkling point highlights: one wandering point per cell, each with its own
// blink rhythm. `r` is the point radius in cell units; it never shrinks below
// ~1.5 px so sparkles stay stable instead of aliasing.
float glitter(vec2 p, float t, float r) {
    vec2 cell = floor(p), f = fract(p);
    float px = fwidth(p.x) * 1.5;
    float rr = max(r, px);
    float energy = r * r / (rr * rr);  // keep brightness constant when widened
    float s = 0.0;
    for (int y = -1; y <= 1; ++y)
        for (int x = -1; x <= 1; ++x) {
            vec2 id = cell + vec2(x, y);
            vec2 h = vec2(hash12(id), hash12(id + 19.19));
            vec2 pt = vec2(x, y) + 0.5 + 0.35 * sin(t * (0.4 + h * 0.6) + h.yx * 6.2832);
            float blink = pow(max(sin(t * (1.5 + 3.0 * h.x) + h.y * 40.0), 0.0), 6.0);
            s += (1.0 - smoothstep(rr * 0.4, rr, length(f - pt))) * blink;
        }
    return s * energy;
}

// Expanding rings from raindrops: one drop per cell per cycle, random phase.
// Returns the surface gradient contribution.
vec2 rainRipples(vec2 p, float t) {
    vec2 cell = floor(p), f = fract(p);
    vec2 grad = vec2(0.0);
    for (int y = -1; y <= 1; ++y)
        for (int x = -1; x <= 1; ++x) {
            vec2 id = cell + vec2(x, y);
            vec2 c = vec2(x, y) + vec2(hash12(id), hash12(id + 5.3)) * 0.8 + 0.1;
            float phase = fract(t * 1.3 + hash12(id + 11.1));
            vec2 d = f - c;
            float r = length(d);
            float front = phase * 0.9;  // ring radius in cell units
            float ring = sin((r - front) * 38.0) * exp(-(r - front) * (r - front) * 120.0) * (1.0 - phase);
            grad += (r > 1e-3 ? d / r : vec2(0.0)) * ring;
        }
    return grad;
}

// Returns (normal.xyz, normalized height) of layered drifting wavelets.
vec4 waterSurface(vec2 p, float t, float shoreDist) {
    vec2 rain = uWeather.x > 0.0
        ? (rainRipples(p * 2.4, t) + rainRipples(p * 3.7 + 9.0, t * 1.17)) * 0.3 * uWeather.x
        : vec2(0.0);
    vec2 grad = vec2(0.0);
    float h = 0.0, amp = 1.0, freq = uWaveScale, norm = 0.0;
    const mat2 rot = mat2(0.8, -0.6, 0.6, 0.8);
    vec2 dir = vec2(1.0, 0.35);
    for (int i = 0; i < 5; ++i) {
        vec3 n = noised(p * freq + dir * t * uWaveSpeed * (1.0 + float(i) * 0.4));
        h += n.x * amp;
        norm += amp;
        grad += n.yz * amp * freq;
        p = rot * p;
        dir = rot * dir;
        amp *= 0.5;
        freq *= 1.85;
    }
    grad *= uWaveStrength * 0.09;
    // Ripples rolling in toward the shore.
    float ring = sin(shoreDist * 9.0 + t * 2.4) * exp(-shoreDist * 1.6) * 0.06 * uWaveStrength;
    grad += vec2(ring) + rain;
    return vec4(normalize(vec3(-grad.x, 1.0, -grad.y)), h / norm);
}

void main() {
    vec3 wp = vWorld;
    vec3 V = normalize(uCamPos.xyz - wp);
    float t = uTime;
    vec2 suv = gl_FragCoord.xy * uScreen.zw;
    float shoreDist = texture(uShoreDist, (wp.xz - uShoreRect.xy) * uShoreRect.zw).r;
    vec4 surf = waterSurface(wp.xz, t, shoreDist);
    vec3 n = surf.xyz;

    // Bottom under this pixel.
    vec3 sceneP = worldFromDepth(suv, texture(uSceneDepth, suv).r);
    float depthV = max(wp.y - sceneP.y, 0.0);

    // Refraction: offset by the surface normal, scaled down in the shallows.
    vec2 ruv = suv + n.xz * uRefraction * clamp(depthV * 2.5, 0.0, 1.0);
    vec3 rP = worldFromDepth(ruv, texture(uSceneDepth, ruv).r);
    if (rP.y > wp.y + 0.01) {  // refracted sample is above water: fall back
        ruv = suv;
        rP = sceneP;
    }
    vec3 refr = texture(uSceneColor, ruv).rgb;
    float rDepth = max(wp.y - rP.y, 0.0);
    float path = length(rP - wp) + rDepth;
    vec3 transmittance = exp(-uAbsorb * path / max(uClarity, 0.05));

    float sunVis = sunVisibility(wp);
    vec3 L = uSunDir.xyz;
    vec3 sunRad = uSunColor.rgb * uSunDir.w;
    vec3 lightIn = sunRad * max(L.y, 0.0) * mix(0.35, 1.0, sunVis) + hemiAmbient(vec3(0, 1, 0));
    vec3 body = mix(uShallow, uDeep, 1.0 - exp(-rDepth * 0.8)) * lightIn * 0.65;
    vec3 under = refr * transmittance + body * (1.0 - transmittance);

    // Reflection: planar reflection texture shares screen mapping with the main view.
    vec3 R = reflect(-V, n);
    vec3 refl = uHasReflection == 1 ? texture(uReflection, suv + n.xz * uReflDistort).rgb : skyColor(R);
    refl *= mix(vec3(1.0), normalize(uShallow + 1e-3) * 1.6, 0.35);  // water tints what it reflects
    float ndv = max(dot(n, V), 0.0);
    float F = 0.02 + 0.98 * pow(1.0 - ndv, 5.0);
    F = mix(F, 1.0, 0.06);  // a hint of reflection even looking straight down
    vec3 col = mix(under, refl, F);

    // Sky-lit wave crests: stylized bright streaks that keep the surface alive
    // even when the sun's reflection is out of view.
    float crest = smoothstep(0.58, 0.78, surf.w);
    col += skyColor(R) * crest * 0.10 * (0.5 + 0.5 * sunVis);

    // Sun glints: tight core + wide sheen (HDR, drives bloom). Shadowed water
    // still catches some sky glints so sparkles don't vanish under trees.
    vec3 H = normalize(L + V);
    float nh = max(dot(n, H), 0.0);
    float specVis = mix(0.15, 1.0, sunVis);
    float spec = pow(nh, 1200.0) * 120.0 + pow(nh, 160.0) * 4.0 + pow(nh, 24.0) * 0.35;
    col += sunRad * spec * uSpecular * 0.05 * specVis;

    // Glitter: a field of twinkling point highlights on wave facets, two scales.
    // Each is a few pixels wide and very bright, so bloom and depth of field
    // turn them into the soft bokeh sparkles typical of HD-2D water.
    float lobe = 0.35 + pow(nh, 8.0) * 3.0;
    float g = glitter(wp.xz * 3.5, t, 0.13) + glitter(wp.xz * 1.6 + 17.0, t * 0.8, 0.085) * 1.8;
    g *= smoothstep(0.3, 0.7, surf.w) * (0.5 + 4.0 * F);
    col += sunRad * g * lobe * uGlitter * 8.0 * specVis * max(L.y, 0.25);

    // Shoreline foam: broken bands drifting along the distance-to-shore field.
    float fn = vnoise(wp.xz * 4.0 + vec2(t * 0.35, -t * 0.25)) * 0.6 + vnoise(wp.xz * 9.0 - t * 0.5) * 0.4;
    float d = shoreDist + (fn - 0.5) * 0.12;
    float edge = 1.0 - smoothstep(0.0, uFoamWidth * 0.3, d);                   // solid contact line
    float bands = sin(d / max(uFoamWidth, 1e-3) * 9.0 - t * 1.8) * 0.5 + 0.5;  // receding wash
    float wash = (1.0 - smoothstep(0.0, uFoamWidth, d)) * smoothstep(0.55, 0.85, bands * 0.6 + fn * 0.5);
    float foam = clamp(max(edge, wash) * uFoamIntensity, 0.0, 1.0);
    vec3 foamCol = vec3(0.92, 0.96, 1.0) * (sunRad * max(L.y, 0.0) * mix(0.4, 1.0, sunVis) + hemiAmbient(vec3(0, 1, 0))) * 0.55;
    col = mix(col, foamCol, foam);

    fragColor = vec4(col, 1.0);
}