import { GLSL3, Vector2 } from "three";

/**
 * Default ray count for the AO baker. Higher = less noise, longer bake.
 * Variance falls as 1/sqrt(NUM_RAYS), so doubling rays halves stdev.
 *
 * @type {number}
 */
const DEFAULT_NUM_RAYS = 64;

/**
 * Lightmap baker for heightfield terrain.
 *
 * Per output texel, traces `NUM_RAYS` cosine-weighted rays into the
 * hemisphere around the surface normal and counts how many are occluded by
 * the heightfield. The output channel is visibility ∈ [0, 1]:
 *   1 = fully lit (no occlusion)
 *   0 = fully occluded
 *
 * Architecture
 * ------------
 *
 *   1. Pre-pass (orchestrator side): a max-reduction mipmap pyramid is built
 *      from the heightmap — see `build_max_height_pyramid`. Each mip stores
 *      per-cell upper bounds on height. The whole pyramid is one DataTexture
 *      with `tex.mipmaps[]`, sampled by `textureLod`.
 *
 *   2. Sampling: directions are drawn from a cosine-weighted hemisphere using
 *      Shirley's concentric disk mapping. Cosine-weighting matches the
 *      Lambertian diffuse assumption — each sample's contribution to the AO
 *      integral has unit weight, and the estimator collapses to a simple
 *      hits/total ratio.
 *
 *   3. Random seeding: a PCG hash (pcg3d on the fragment coordinate) seeds an
 *      independent stream per output texel. Subsequent rays draw from `pcg`
 *      on the running state. White noise — no Halton/blue-noise sequences —
 *      keeps the implementation simple and the noise spectrum unbiased; a
 *      bilateral denoise pass downstream cleans up the speckle.
 *
 *   4. Tracing: Hi-Z traversal in world space. The ray's XZ projection walks
 *      texture-space cells; per cell, the max-height at the current mip is
 *      compared against the ray's min height across the cell. Above max:
 *      advance to the next cell and try to coarsen the mip. Below max:
 *      refine the mip; at mip 0, record a hit. Bounded by MAX_HIZ_ITER per
 *      ray as a safety net for pathological convergence.
 *
 * World space
 * -----------
 *
 *   Convention: X and Z are horizontal (UV-aligned), Y is up. The heightmap's
 *   UV maps linearly to world (x, z) ∈ [0, worldSize.x] × [0, worldSize.y],
 *   and the heightmap value at any UV is the world Y coordinate of the
 *   surface. Surface normals are recomputed in-shader by central differences
 *   on the heightmap, removing any dependency on an external normal map.
 *
 * Inputs (uniforms — plumbing, not user-facing)
 * ---------------------------------------------
 *
 *   heightMap     R32F texture with a max-reduction mipmap chain
 *   worldSize     XZ extent of the terrain in world units
 *   heightMapSize Resolution of heightMap mip 0 in texels
 *   maxMipLevel   Index of the coarsest mip in heightMap
 *
 * The single user-facing parameter is `NUM_RAYS`, configured by the factory
 * argument.
 *
 * @param {Object} [opts]
 * @param {number} [opts.numRays=64] number of rays per output texel
 */
function AmbientOcclusionShader({ numRays = DEFAULT_NUM_RAYS } = {}) {
    return {
        uniforms: {
            heightMap:     { value: null },
            worldSize:     { value: new Vector2(1, 1) },
            heightMapSize: { value: new Vector2(1, 1) },
            maxMipLevel:   { value: 0 }
        },

        defines: {
            NUM_RAYS: numRays | 0
        },

        glslVersion: GLSL3,

        vertexShader: `
            out vec2 vUv;

            void main() {
                vUv = uv;
                gl_Position = vec4((uv - 0.5) * 2.0, 0.0, 1.0);
            }
        `,

        fragmentShader: `
            // three.js auto-prefixes float/int precision; declare highp for
            // the heightmap sampler so R32F reads keep full precision.
            //
            // The pyramid is stored as a DataTexture2DArray with one mip per
            // layer (all at full mip-0 resolution; smaller mips are nearest-
            // upscaled). This sidesteps a three.js r136 upload bug for
            // DataTexture custom mipmap chains — see build_max_height_pyramid.
            precision highp sampler2DArray;

            uniform sampler2DArray heightMap;
            uniform vec2 worldSize;
            uniform vec2 heightMapSize;
            uniform int maxMipLevel;

            in vec2 vUv;
            out vec4 fragColor;

            const float PI = 3.14159265358979;

            // Safety cap on Hi-Z iterations per ray. log2 of typical heightmap
            // dimensions × a handful of refine/coarsen rounds per cell is
            // usually under 80; 256 gives generous headroom.
            const int MAX_HIZ_ITER = 256;

            // ============================================================
            //  PCG random number generator
            // ============================================================
            //  pcg is the standard 32-bit scrambler from O'Neill's family.
            //  pcg3d is the 3-vector variant used purely for seeding so the
            //  spatial pattern of per-texel seeds is uncorrelated.

            uint rng_state;

            uint pcg(uint v) {
                uint state = v * 747796405u + 2891336453u;
                uint word = ((state >> ((state >> 28u) + 4u)) ^ state) * 277803737u;
                return (word >> 22u) ^ word;
            }

            uvec3 pcg3d(uvec3 v) {
                v = v * 1664525u + 1013904223u;
                v.x += v.y * v.z;
                v.y += v.z * v.x;
                v.z += v.x * v.y;
                v ^= (v >> uvec3(16u));
                v.x += v.y * v.z;
                v.y += v.z * v.x;
                v.z += v.x * v.y;
                return v;
            }

            // Map a uint to [0, 1) by stuffing 23 bits into a float's mantissa
            // and subtracting the implicit leading 1. Faster than divide and
            // produces a uniform distribution.
            float hash_to_float01(uint h) {
                return uintBitsToFloat(0x3f800000u | (h >> 9u)) - 1.0;
            }

            uint random_uint() {
                rng_state = pcg(rng_state);
                return rng_state;
            }

            float random() {
                return hash_to_float01(random_uint());
            }

            void random_initialize(uvec3 invocation_id, uvec3 seed) {
                uvec3 mixed = pcg3d(invocation_id + seed * 37u);
                rng_state = mixed.x ^ mixed.y ^ mixed.z;
            }

            // ============================================================
            //  Cosine-weighted hemisphere sampling
            // ============================================================
            //  Shirley's concentric disk → cosine hemisphere via z = sqrt(1 - r²).
            //  The orthonormal basis from the normal uses Frisvad's branchless
            //  construction with z-flip for numerical stability near n.z = -1.

            vec2 sample_uniform_disk_concentric(vec2 uv) {
                vec2 offset = 2.0 * uv - vec2(1.0);
                if (offset.x == 0.0 && offset.y == 0.0) {
                    return vec2(0.0);
                }
                float r;
                float theta;
                const float PI_4 = PI / 4.0;
                const float PI_2 = PI / 2.0;
                if (abs(offset.x) > abs(offset.y)) {
                    r = offset.x;
                    theta = PI_4 * (offset.y / offset.x);
                } else {
                    r = offset.y;
                    theta = PI_2 - PI_4 * (offset.x / offset.y);
                }
                return r * vec2(cos(theta), sin(theta));
            }

            vec3 sample_cosine_weighted_hemisphere(vec2 uv) {
                vec2 d = sample_uniform_disk_concentric(uv);
                float z = sqrt(max(0.0, 1.0 - d.x * d.x - d.y * d.y));
                return vec3(d.x, d.y, z);
            }

            mat3 build_orthonormal_matrix_n(vec3 n) {
                vec3 T;
                vec3 B;
                if (n.z < 0.0) {
                    float a = 1.0 / (1.0 - n.z);
                    float b = n.x * n.y * a;
                    T = vec3(1.0 - n.x * n.x * a, -b, n.x);
                    B = vec3(b, n.y * n.y * a - 1.0, -n.y);
                } else {
                    float a = 1.0 / (1.0 + n.z);
                    float b = -n.x * n.y * a;
                    T = vec3(1.0 - n.x * n.x * a, b, -n.x);
                    B = vec3(b, 1.0 - n.y * n.y * a, -n.y);
                }
                return mat3(T, B, n);
            }

            vec3 cosine_weighted_world_direction(vec2 uv, vec3 n) {
                vec3 local = sample_cosine_weighted_hemisphere(uv);
                return normalize(build_orthonormal_matrix_n(n) * local);
            }

            // ============================================================
            //  Heightmap utilities
            // ============================================================

            // Sample mip-0 (the original heightmap) at uv. Layer 0 of the
            // 2D array holds the un-reduced heightmap data.
            float sample_height(vec2 uv) {
                return texture(heightMap, vec3(uv, 0.0)).r;
            }

            // Sample max-height at a specific Hi-Z level. Mip k lives in
            // layer k, nearest-upscaled to mip-0 resolution. NearestFilter
            // on the array snaps the layer to integer 'mip' and returns
            // the underlying mip-k texel containing uv.
            float sample_max_height(vec2 uv, int mip) {
                return texture(heightMap, vec3(uv, float(mip))).r;
            }

            // Central-difference normal in world space. Faster than building
            // a normal map up front and stays self-consistent with the
            // heightmap used for tracing (no convention mismatch possible).
            vec3 compute_world_normal(vec2 uv) {
                vec2 texel_step  = 1.0 / heightMapSize;
                vec2 texel_world = worldSize / heightMapSize;

                float h_l = sample_height(uv - vec2(texel_step.x, 0.0));
                float h_r = sample_height(uv + vec2(texel_step.x, 0.0));
                float h_d = sample_height(uv - vec2(0.0, texel_step.y));
                float h_u = sample_height(uv + vec2(0.0, texel_step.y));

                float dHdx = (h_r - h_l) / (2.0 * texel_world.x);
                float dHdz = (h_u - h_d) / (2.0 * texel_world.y);

                return normalize(vec3(-dHdx, 1.0, -dHdz));
            }

            // ============================================================
            //  Hi-Z ray traversal
            // ============================================================
            //  Trace one ray against the max-height pyramid. The ray's XZ
            //  projection steps through texture-space cells; height is tracked
            //  separately along the ray parameter t (world units).
            //
            //  Returns true iff the ray is occluded by the heightfield before
            //  exiting the terrain extent.

            bool trace_ray_hiz(vec3 origin_w, vec3 dir_w) {
                vec2 inv_world_size = 1.0 / worldSize;
                vec2 origin_tex = origin_w.xz * inv_world_size;
                vec2 dir_tex    = dir_w.xz   * inv_world_size;

                // Reciprocal of dir_tex used to compute "world distance to
                // next cell boundary". A zero component means the ray is
                // axis-aligned in tex space — it never crosses that axis's
                // boundaries — so route to a large finite value with sign
                // matching step(0, dir) (= +1 at the zero crossing).
                vec2 inv_dir_tex = vec2(
                    dir_tex.x == 0.0 ? 1e20 : 1.0 / dir_tex.x,
                    dir_tex.y == 0.0 ? 1e20 : 1.0 / dir_tex.y
                );

                // Step past the cell boundary by a tiny world-space epsilon
                // so the next iteration sees the new cell, not the one we
                // just exited.
                const float T_EPS = 1e-4;

                int   mip = 0;
                float t   = 0.0;

                for (int iter = 0; iter < MAX_HIZ_ITER; iter++) {
                    vec2 tex_pos = origin_tex + t * dir_tex;

                    // Exited the terrain extent — ray escaped, no occlusion
                    if (any(lessThan(tex_pos, vec2(0.0))) ||
                        any(greaterThan(tex_pos, vec2(1.0)))) {
                        return false;
                    }

                    // Cell extents at the current mip
                    float mip_scale   = exp2(float(mip));
                    vec2  cell_size   = vec2(mip_scale) / heightMapSize;
                    vec2  cell_origin = floor(tex_pos / cell_size) * cell_size;
                    vec2  cell_center = cell_origin + 0.5 * cell_size;

                    // World distance along the ray to the next cell boundary
                    vec2 next_boundary = cell_origin + step(0.0, dir_tex) * cell_size;
                    vec2 t_to_boundary = max((next_boundary - tex_pos) * inv_dir_tex, vec2(0.0));
                    float t_exit = min(t_to_boundary.x, t_to_boundary.y) + T_EPS;

                    // Max occluder height in this cell at this mip
                    float max_h = sample_max_height(cell_center, mip);

                    // Ray height is linear in t, so the minimum across the
                    // cell is at one of the endpoints
                    float h_now  = origin_w.y + t * dir_w.y;
                    float h_exit = origin_w.y + (t + t_exit) * dir_w.y;
                    float h_min  = min(h_now, h_exit);

                    if (h_min <= max_h) {
                        // Cell could contain an occluder
                        if (mip == 0) {
                            // Finest mip: confirmed hit
                            return true;
                        }
                        // Refine into smaller cells without advancing t
                        mip = mip - 1;
                    } else {
                        // Cell is empty above the ray — skip past it and try
                        // to coarsen so the next iteration takes a bigger step
                        t   = t + t_exit;
                        mip = min(maxMipLevel, mip + 1);
                    }
                }

                // Iteration cap reached — choose the unbiased side and return
                // "not occluded". Visibility errors on the bright side are
                // less visible than spurious dark pixels.
                return false;
            }

            // ============================================================
            //  Main
            // ============================================================

            void main() {
                // Seed an independent PCG stream per output texel
                uvec2 frag = uvec2(gl_FragCoord.xy);
                random_initialize(uvec3(frag, 0u), uvec3(0u));

                // World-space origin: lift the texel out of the heightmap
                vec2  world_xz = vUv * worldSize;
                float h0       = sample_height(vUv);
                vec3  origin_w = vec3(world_xz.x, h0, world_xz.y);

                vec3 normal_w = compute_world_normal(vUv);

                // Push origin a fraction of a texel along the normal so the
                // first Hi-Z cell at mip 0 doesn't trivially self-occlude
                vec2  texel_world = worldSize / heightMapSize;
                float origin_bias = 0.5 * min(texel_world.x, texel_world.y);
                origin_w += normal_w * origin_bias;

                int hits = 0;

                for (int i = 0; i < NUM_RAYS; i++) {
                    vec2 r2  = vec2(random(), random());
                    vec3 dir = cosine_weighted_world_direction(r2, normal_w);

                    if (trace_ray_hiz(origin_w, dir)) {
                        hits++;
                    }
                }

                // Cosine-weighted sampling makes each ray's contribution
                // unit-weighted in the AO integral, so the estimator is just
                // the fraction of escaped rays.
                float occlusion  = float(hits) / float(NUM_RAYS);
                float visibility = 1.0 - occlusion;

                fragColor = vec4(visibility, 0.0, 0.0, 1.0);
            }
        `
    };
}

export default AmbientOcclusionShader;
