// Screen-space ambient occlusion, evaluated from the camera linear-depth
// prepass.
//
// Reconstructs each pixel's view-space position and face normal from depth,
// then averages normalized angular obscurance over a world-space sampling
// radius. The output is a single occlusion factor in the red channel: 1 is
// unoccluded, 0 is fully occluded. A bilateral blur pass cleans up the
// per-pixel noise afterwards.
//
// The depth hierarchy and sampling follow McGuire et al., "Scalable Ambient
// Obscurance" (2012),
// https://research.nvidia.com/publication/scalable-ambient-obscurance.

uniform sampler2D linear_depth;

// Downsampled copies of the depth for the SAO mip chain (half, quarter, and
// eighth resolution). When the chain is disabled these are all bound to
// linear_depth and params2.y is 1, so the sample loop only ever reads level 0.
uniform sampler2D depth_mip1;
uniform sampler2D depth_mip2;
uniform sampler2D depth_mip3;

uniform SsaoInfo {
  // x, y: occlusion target size in pixels. z, w: its reciprocal.
  vec4 viewport;
  // x: tan(fovX / 2). y: tan(fovY / 2). z: near plane. w: far plane.
  vec4 proj;
  // x: radius (world units). y: bias (world units). z: obscurance intensity.
  // w: projection scale (pixels per world unit at depth 1).
  vec4 params;
  // x: sample count. y: mip level count. z: minimum horizon sine term.
  // w: final visibility power.
  vec4 params2;
  // x: immediate-neighbour detail intensity. yzw: unused.
  vec4 params3;
}
ssao;

in vec2 v_uv;
out vec4 frag_color;

// Static loop bound so the sample loop is constant-bounded for the GLES
// 1.00 shader output (a uniform-bounded loop is rejected by conformant ES
// drivers); the dynamic sample count breaks out early.
#define MAX_SSAO_SAMPLES 32

const float kPi = 3.14159265359;
const float kNumericEpsilon = 0.0001;

#include <interleaved_gradient_noise.glsl>

// Reconstructs a view-space position from a depth-buffer UV. Camera space
// places the eye at the origin looking down +Z (the convention the depth
// prepass writes), so the stored planar depth is the view-space Z and the
// X/Y follow from the projection tangents.
// Fetches the view-space depth at [uv] from level [level] of the depth chain.
float DepthAtLevel(vec2 uv, int level) {
  if (level <= 0) return texture(linear_depth, uv).r;
  if (level == 1) return texture(depth_mip1, uv).r;
  if (level == 2) return texture(depth_mip2, uv).r;
  return texture(depth_mip3, uv).r;
}

vec3 ViewPositionAt(vec2 uv, int level) {
  float z = DepthAtLevel(uv, level);
  // NDC from the full-screen UV (V runs downward in the UV).
  vec2 ndc = vec2(2.0 * uv.x - 1.0, 1.0 - 2.0 * uv.y);
  return vec3(ndc.x * z * ssao.proj.x, ndc.y * z * ssao.proj.y, z);
}

vec3 ViewPositionBase(ivec2 coord) {
  ivec2 size = textureSize(linear_depth, 0);
  coord = clamp(coord, ivec2(0), size - ivec2(1));
  float z = texelFetch(linear_depth, coord, 0).r;
  vec2 uv = (vec2(coord) + vec2(0.5)) / vec2(size);
  vec2 ndc = vec2(2.0 * uv.x - 1.0, 1.0 - 2.0 * uv.y);
  return vec3(ndc.x * z * ssao.proj.x, ndc.y * z * ssao.proj.y, z);
}

// Selects the smoother one-sided derivative on each axis. This avoids joining
// foreground and background positions at silhouettes, while retaining fp32
// position differences across broad sloped surfaces.
vec3 ReconstructNormal(vec2 uv, vec3 center) {
  ivec2 size = textureSize(linear_depth, 0);
  ivec2 p = clamp(ivec2(uv * vec2(size)), ivec2(0), size - ivec2(1));
  vec3 center_sample = ViewPositionBase(p);
  vec3 left = ViewPositionBase(p - ivec2(1, 0));
  vec3 right = ViewPositionBase(p + ivec2(1, 0));
  vec3 down = ViewPositionBase(p - ivec2(0, 1));
  vec3 up = ViewPositionBase(p + ivec2(0, 1));
  vec3 dx;
  if (p.x == 0) {
    dx = right - center_sample;
  } else if (p.x == size.x - 1) {
    dx = center_sample - left;
  } else {
    dx = length(center_sample - left) < length(right - center_sample)
        ? center_sample - left
        : right - center_sample;
  }
  vec3 dy;
  if (p.y == 0) {
    dy = up - center_sample;
  } else if (p.y == size.y - 1) {
    dy = center_sample - down;
  } else {
    dy = length(center_sample - down) < length(up - center_sample)
        ? center_sample - down
        : up - center_sample;
  }
  vec3 normal = normalize(cross(dx, dy));
  return dot(normal, center) > 0.0 ? -normal : normal;
}

float SampleObscurance(vec3 normal, vec3 delta, float falloff_scale,
                       float horizon) {
  float distance2 = dot(delta, delta);
  float normal_dot_delta =
      dot(normal, delta) * inversesqrt(max(distance2, kNumericEpsilon));
  float falloff = max(0.0, 1.0 + distance2 * falloff_scale);
  return max(0.0, normal_dot_delta - horizon) * falloff;
}

void main() {
  float radius = ssao.params.x;
  float bias = ssao.params.y;
  float intensity = ssao.params.z;
  float proj_scale = ssao.params.w;
  int sample_count = int(ssao.params2.x);
  int mip_levels = int(ssao.params2.y);
  float horizon_angle = ssao.params2.z;
  float power = ssao.params2.w;
  float detail = ssao.params3.x;
  float far = ssao.proj.w;

  vec3 origin = ViewPositionAt(v_uv, 0);

  // Background texels (no geometry) are unoccluded.
  if (origin.z >= far) {
    frag_color = vec4(1.0, 1.0, 1.0, 1.0);
    return;
  }

  vec3 normal = ReconstructNormal(v_uv, origin);

  // Screen-space disk radius: the world radius projected to this depth.
  // Keep the full radius at viewport boundaries. Depth sampling clamps there,
  // so the same normalized kernel remains stable without a separate output
  // fade or a boundary-dependent change in contact strength.
  float screen_radius = 0.85 * proj_scale * radius / origin.z;

  // A non-repeating pixel hash avoids the screen-aligned 4x4 phase pattern
  // that survives filtering as horizontal and vertical bands.
  float noise = InterleavedGradientNoise(gl_FragCoord.xy);
  float rotation = noise * 2.0 * kPi;
  float spiral_turns = sample_count <= 8 ? 3.0
      : (sample_count <= 12 ? 5.0 : (sample_count <= 20 ? 7.0 : 14.0));

  float radius2 = radius * radius;
  float falloff_scale = -1.0 / max(radius2, kNumericEpsilon);
  float horizon = horizon_angle + max(bias, 0.0) / max(radius, kNumericEpsilon);
  float sum = 0.0;
  float weight_sum = 0.0;
  for (int i = 0; i < MAX_SSAO_SAMPLES; i++) {
    if (i >= sample_count) {
      break;
    }
    // The linear radial progression and prime-ish turn counts follow the SAO
    // reference kernel. Unlike sqrt spacing, this retains enough close taps to
    // resolve contact occlusion without amplifying distant depth variation.
    float sample_radius = (float(i) + 0.5) / float(sample_count);
    float theta = sample_radius * spiral_turns * 2.0 * kPi + rotation;
    float pixel_radius = sample_radius * screen_radius;
    // Fade subpixel support continuously instead of dropping taps as the
    // projected kernel shrinks. The one-pixel floor keeps depth taps distinct
    // enough for nearest sampling while the support removes their influence.
    float sample_support = smoothstep(0.0, 1.0, pixel_radius);
    pixel_radius = max(1.0, pixel_radius);
    vec2 offset = vec2(cos(theta), sin(theta)) * pixel_radius;
    vec2 sample_uv = v_uv + offset * ssao.viewport.zw;

    // Read the depth from a coarser level as the sample gets further away, so
    // far taps read pre-reduced depth (cache-friendly) and, with a full-
    // resolution level 0, near geometry is not contaminated by the surface
    // behind it where the projection compresses depth. Level 0 when the chain
    // is disabled (mip_levels <= 1).
    int level = 0;
    if (mip_levels > 1) {
      level = clamp(
          int(floor(log2(max(1.0, pixel_radius)))) - 3, 0, mip_levels - 1);
    }
    vec3 q = ViewPositionAt(sample_uv, level);
    vec3 v = q - origin;
    float reduce = clamp(
        max(0.0, -v.z) * (-1.0 / max(radius, kNumericEpsilon)) + 2.0,
        0.0, 1.0);
    float weight = 0.6 * reduce + 0.4;
    sum += SampleObscurance(normal, v, falloff_scale, horizon) * weight *
        sample_support;
    weight_sum += weight;
  }

  // Immediate neighbours recover tight contact detail before the wide kernel
  // is filtered. Their falloff is four times faster, matching the main radius
  // normalisation used by the reference ASSAO implementation.
  ivec2 center_coord = ivec2(v_uv * vec2(textureSize(linear_depth, 0)));
  float detail_sum = 0.0;
  detail_sum += SampleObscurance(
      normal, ViewPositionBase(center_coord + ivec2(-1, 0)) - origin,
      falloff_scale * 4.0, horizon);
  detail_sum += SampleObscurance(
      normal, ViewPositionBase(center_coord + ivec2(1, 0)) - origin,
      falloff_scale * 4.0, horizon);
  detail_sum += SampleObscurance(
      normal, ViewPositionBase(center_coord + ivec2(0, -1)) - origin,
      falloff_scale * 4.0, horizon);
  detail_sum += SampleObscurance(
      normal, ViewPositionBase(center_coord + ivec2(0, 1)) - origin,
      falloff_scale * 4.0, horizon);

  float wide_obscurance = weight_sum > 0.0 ? sum / weight_sum : 0.0;
  // Both terms are means. `detail` is therefore the additive strength of the
  // four-neighbour mean before `intensity` scales their combined obscurance.
  float obscurance = wide_obscurance + detail * detail_sum * 0.25;
  obscurance = min(intensity * obscurance, 0.98);
  float ao = pow(clamp(1.0 - obscurance, 0.0, 1.0), power);

  frag_color = vec4(ao, ao, ao, 1.0);
}
