import { ClampToEdgeWrapping, DataTexture, RGBAFormat, UnsignedByteType, Vector2 } from 'three';
import ImageFilter from '../../graphics/filter/ImageFilter.js';
import AmbientOcclusionShader from '../../graphics/shaders/AmbientOcclusionShader.js';
import { DenoiseShader } from "../../graphics/shaders/DenoiseShader.js";
import { Sampler2D } from '../../graphics/texture/sampler/Sampler2D.js';
import { build_max_height_pyramid } from "../../graphics/util/build_max_height_pyramid.js";

/**
 * Wrap a Uint8 buffer as an RGBA8 texture suitable for the denoise pass.
 *
 * @param {Uint8Array|Uint8ClampedArray} data
 * @param {number} width
 * @param {number} height
 * @returns {DataTexture}
 */
function filter_result_to_texture(data, width, height) {
    const result = new DataTexture();
    result.format = RGBAFormat;
    result.type = UnsignedByteType;
    result.flipY = false;
    result.image = { data, width, height };
    result.wrapS = ClampToEdgeWrapping;
    result.wrapT = ClampToEdgeWrapping;
    result.repeat.set(1, 1);
    result.needsUpdate = true;
    result.anisotropy = 4;
    return result;
}

/**
 * Bake an ambient occlusion / visibility map for a heightfield terrain.
 *
 * Pipeline:
 *   1. Build a max-reduction Hi-Z pyramid from the heightmap (pre-pass).
 *   2. Run {@link AmbientOcclusionShader} which traces `rayCount`
 *      cosine-weighted rays per output texel against the pyramid.
 *   3. Bilateral denoise to clean up Monte Carlo speckle.
 *
 * The output is a single-channel visibility map (1 = fully lit, 0 = fully
 * occluded) packed into a {@link Sampler2D}.
 *
 * @param {WebGLRenderer} renderer
 * @param {Sampler2D} heightMap        single-channel Float32 heightmap
 * @param {Vector2} resultSize         output AO resolution in texels
 * @param {Vector2} [worldSize]        terrain XZ extent in world units;
 *   defaults to heightmap pixel dimensions (= 1 world unit per heightmap texel)
 * @param {number} [rayCount=64]       Monte Carlo samples per output texel
 * @returns {Sampler2D}                single-channel Uint8 visibility map
 */
function heightMap2AOMap(
    renderer,
    heightMap,
    resultSize,
    worldSize = new Vector2(heightMap.width, heightMap.height),
    rayCount = 64
) {
    const width = resultSize.x;
    const height = resultSize.y;

    // Pre-pass: max-height pyramid for Hi-Z ray traversal.
    // The pyramid is a DataTexture2DArray; each Hi-Z mip is one layer.
    const pyramid_texture = build_max_height_pyramid(heightMap);
    const max_mip_level = pyramid_texture.userData.mipCount - 1;

    // AO bake
    const shader_ao = AmbientOcclusionShader({ numRays: rayCount });
    shader_ao.uniforms.heightMap.value = pyramid_texture;
    shader_ao.uniforms.worldSize.value.set(worldSize.x, worldSize.y);
    shader_ao.uniforms.heightMapSize.value.set(heightMap.width, heightMap.height);
    shader_ao.uniforms.maxMipLevel.value = max_mip_level;

    const raw_ao = ImageFilter(renderer, width, height, shader_ao);

    // Denoise: bilateral filter smooths Monte Carlo noise while preserving
    // visibility edges (cliff bases, ridges)
    const shader_denoise = new DenoiseShader(2);
    shader_denoise.uniforms.source.value = filter_result_to_texture(raw_ao.array, width, height);
    shader_denoise.uniforms.pixelOffset.value.set(1 / width, 1 / height);

    const smooth_ao = ImageFilter(renderer, width, height, shader_denoise);

    // Pack red channel back into a single-channel Uint8 sampler — that's what
    // the terrain material consumes as aoMap
    const result = new Sampler2D(new Uint8ClampedArray(width * height), 1, width, height);
    const size = width * height;
    for (let i = 0; i < size; i++) {
        result.data[i] = smooth_ao.array[i * 4];
    }

    // Free the upload-only pyramid texture
    pyramid_texture.dispose();

    return result;
}

export default heightMap2AOMap;
