import { ClampToEdgeWrapping, DataTexture2DArray, FloatType, NearestFilter, RedFormat } from "three";

/**
 * Build a max-reduction Hi-Z height pyramid from a single-channel float
 * heightmap sampler.
 *
 * Each pyramid level k holds, per texel, the maximum height across the 2×2
 * block of level k-1 it covers. A ray trace can then skip whole cells
 * safely: if the ray's height is strictly above the cell's max-height at
 * any level, the cell can be advanced past without checking any finer level.
 *
 * Storage layout — **DataTexture2DArray, one mip per layer**, all layers at
 * full mip-0 resolution. Smaller mips are nearest-upscaled into their layer.
 * The shader samples mip k via `texture(heightMap, vec3(uv, float(k)))`.
 *
 * This sidesteps the three.js r136 bug in `WebGLTextures.uploadTexture`
 * where DataTexture custom mipmap chains are uploaded with `texSubImage2D`
 * called with level hardcoded to 0 — every "mip" smashes the previous one
 * into level 0 and the actual mip levels 1..N are left uninitialized. The
 * 2DArray upload path uses one `texSubImage3D` call to level 0 covering
 * all layers, so it works correctly.
 *
 * Memory cost is (mip_count × W × H × 4) bytes versus (≈ 4/3 × W × H × 4)
 * for a true mip pyramid — about a 3× overhead, traded for correctness.
 *
 * Non-power-of-two heightmaps are handled by halving with floor and stopping
 * once both dimensions reach 1.
 *
 * @param {Sampler2D} sampler  single-channel Float32 heightmap (itemSize=1)
 * @returns {DataTexture2DArray}  R32F array texture; depth = mip count
 */
export function build_max_height_pyramid(sampler) {
    const w0 = sampler.width;
    const h0 = sampler.height;

    // Step 1: build the conventional mip pyramid in JS
    /** @type {Array<{data: Float32Array, width: number, height: number}>} */
    const mips = [];
    mips.push({ data: sampler.data, width: w0, height: h0 });

    {
        let prevW = w0;
        let prevH = h0;
        let prevData = sampler.data;

        while (prevW > 1 || prevH > 1) {
            const curW = Math.max(1, prevW >> 1);
            const curH = Math.max(1, prevH >> 1);
            const curData = new Float32Array(curW * curH);

            const lastSrcX = prevW - 1;
            const lastSrcY = prevH - 1;

            for (let y = 0; y < curH; y++) {
                const sy0 = y << 1;
                const sy1 = Math.min(sy0 + 1, lastSrcY);
                const row0 = sy0 * prevW;
                const row1 = sy1 * prevW;

                for (let x = 0; x < curW; x++) {
                    const sx0 = x << 1;
                    const sx1 = Math.min(sx0 + 1, lastSrcX);

                    const a = prevData[row0 + sx0];
                    const b = prevData[row0 + sx1];
                    const c = prevData[row1 + sx0];
                    const d = prevData[row1 + sx1];

                    let m = a > b ? a : b;
                    if (c > m) m = c;
                    if (d > m) m = d;

                    curData[y * curW + x] = m;
                }
            }

            mips.push({ data: curData, width: curW, height: curH });

            prevW = curW;
            prevH = curH;
            prevData = curData;
        }
    }

    // Step 2: pack each mip nearest-upscaled to w0 × h0 as a layer of the
    // array texture. Layer k holds mip k blown up; sampling with NearestFilter
    // at any uv returns mip k's value at that uv exactly.
    const depth = mips.length;
    const layerSize = w0 * h0;
    const layered = new Float32Array(depth * layerSize);

    for (let k = 0; k < depth; k++) {
        const mip = mips[k];
        const mw = mip.width;
        const mh = mip.height;
        const mdata = mip.data;

        const layerOffset = k * layerSize;

        if (mw === w0 && mh === h0) {
            // Level 0 — direct copy
            layered.set(mdata, layerOffset);
            continue;
        }

        // Nearest-upscale: dst (x, y) → src (floor(x * mw / w0), floor(y * mh / h0))
        // Works for both POT and NPO2 reductions.
        for (let y = 0; y < h0; y++) {
            const sy = Math.min(mh - 1, ((y * mh) / h0) | 0);
            const srcRow = sy * mw;
            const dstRow = layerOffset + y * w0;

            for (let x = 0; x < w0; x++) {
                const sx = Math.min(mw - 1, ((x * mw) / w0) | 0);
                layered[dstRow + x] = mdata[srcRow + sx];
            }
        }
    }

    const tex = new DataTexture2DArray(layered, w0, h0, depth);
    tex.format = RedFormat;
    tex.type = FloatType;
    tex.internalFormat = 'R32F';
    tex.minFilter = NearestFilter;
    tex.magFilter = NearestFilter;
    tex.wrapS = ClampToEdgeWrapping;
    tex.wrapT = ClampToEdgeWrapping;
    tex.generateMipmaps = false;
    tex.flipY = false;
    tex.needsUpdate = true;

    // Expose the depth so callers can plumb `maxMipLevel` to the shader
    // without re-deriving it
    tex.userData.mipCount = depth;

    return tex;
}
