import {
    GLSL3,
    LinearFilter,
    NoBlending,
    RedFormat,
    ShaderMaterial,
    UniformsUtils,
    UnsignedByteType,
    WebGLMultisampleRenderTarget
} from "three";
import { EnginePlugin } from "../../../../../plugin/EnginePlugin.js";
import { CompositingStages } from "../../../../composit/CompositingStages.js";
import { CompositLayer } from "../../../../composit/CompositLayer.js";
import { StandardFrameBuffers } from "../../../../StandardFrameBuffers.js";
import { BlendingType } from "../../../../texture/sampler/BlendingType.js";
import { renderScreenSpace } from "../../../utils/renderScreenSpace.js";
import { generateHilbertNoiseTexture } from "./generateHilbertNoiseTexture.js";
import { SAOShader } from "./SAOShader.js";
import { SAOUpscaleShader } from "./SAOUpscaleShader.js";

export class AmbientOcclusionPostProcessEffect extends EnginePlugin {
    constructor() {
        super();

        this.id = "ambient-occlusion-post-process";

        const defines = Object.assign({}, SAOShader.defines);

        this.__material = new ShaderMaterial({
            defines: defines,
            fragmentShader: SAOShader.fragmentShader,
            vertexShader: SAOShader.vertexShader,
            uniforms: UniformsUtils.clone(SAOShader.uniforms),
            glslVersion: GLSL3,
            depthWrite: false,
            depthTest: false
        });

        // AO parameters
        const uniforms = this.__material.uniforms;

        // artistic darkening control; 1.0 == physically plausible
        uniforms.intensity.value = 1.0;

        // sampling hemisphere radius (also the max occluder distance), in world units (scene dependent)
        uniforms.kernelRadius.value = 0.5;

        this.__material.blending = NoBlending;

        this.__material.extensions.derivatives = true;
        // should depth be unpacked?
        defines.DEPTH_PACKING = 0;
        // is normal buffer available?
        defines.NORMAL_TEXTURE = 0;

        // assume that the camera is a perspective camera
        defines.PERSPECTIVE_CAMERA = 1;

        // More samples = better quality
        defines.NUM_SAMPLES = 24;

        /**
         * Should normal buffer be used or not? Using normal buffer might be significantly more expensive, as this would require normal buffer to be built every frame. If there are other users of the buffer - this is not an issue
         * @type {boolean}
         * @private
         */
        this.__use_normals = false;

        /**
         * Resolution scale at which AO is to be rendered. Higher scale will lead to sharper and better looking results
         * @type {number}
         * @private
         */
        this.__resolution_scale = 0.5;


        /**
         *
         * @type {CompositLayer}
         * @private
         */
        this.__composit_layer = null;

        /**
         *
         * @type {Camera|null}
         * @private
         */
        this.__render_camera = null;

        this.__render_target = new WebGLMultisampleRenderTarget(1, 1, {
            minFilter: LinearFilter,
            magFilter: LinearFilter,

            format: RedFormat,
            type: UnsignedByteType,
            internalFormat: 'R8',
            depthBuffer: false,
            stencilBuffer: false,
            generateMipmaps: false,
        });

        /**
         *
         * @type {FrameBuffer[]}
         * @private
         */
        this.__used_buffers = [];

        /**
         * Edge-aware upscale + denoise: a single non-separable pass that resolves the half-res AO
         * into the full-res composit layer using plane-distance weights.
         * @type {ShaderMaterial}
         * @private
         */
        this.__upscale_material = new ShaderMaterial({
            defines: Object.assign({}, SAOUpscaleShader.defines),
            fragmentShader: SAOUpscaleShader.fragmentShader,
            vertexShader: SAOUpscaleShader.vertexShader,
            uniforms: UniformsUtils.clone(SAOUpscaleShader.uniforms),
            glslVersion: GLSL3,
            depthWrite: false,
            depthTest: false,
            blending: NoBlending
        });
        this.__upscale_material.extensions.derivatives = true;
        // depth is stored unpacked, matching the AO material
        this.__upscale_material.defines.DEPTH_PACKING = 0;

        /**
         * Hilbert-index permutation LUT, shared by the AO sampling and upscale dithers.
         * @type {DataTexture}
         * @private
         */
        this.__hilbert_texture = generateHilbertNoiseTexture(64);
        this.__material.uniforms.tHilbert.value = this.__hilbert_texture;
        this.__upscale_material.uniforms.tHilbert.value = this.__hilbert_texture;
    }

    /**
     * Higher value will make AO appear darker, lower values will make it lighter
     * @param {number} v
     */
    set intensity(v) {
        this.__material.uniforms.intensity.value = v;
    }

    /**
     *
     * @return {number}
     */
    get intensity() {
        return this.__material.uniforms.intensity.value;
    }

    /**
     *
     * @param {Camera} camera
     * @private
     */
    __setCamera(camera) {
        this.__render_camera = camera;

        const desired = camera.isPerspectiveCamera === true ? 1 : 0;

        // the AO and upscale materials both reconstruct view-Z, so both track the projection type
        const materials = [this.__material, this.__upscale_material];
        for (const material of materials) {
            if (material.defines.PERSPECTIVE_CAMERA !== desired) {
                material.defines.PERSPECTIVE_CAMERA = desired;

                // material has changed and will need to be re-compiled
                material.needsUpdate = true;
            }
        }
    }

    finalize() {
        super.finalize();
    }

    __prepare_uniforms() {

        const material = this.__material;

        const uniforms = material.uniforms;

        const camera = this.__render_camera;

        // AO renders at half-res but samples the full-res depth buffer, so `size` is the full (depth)
        // resolution -- all texel math then lands on depth texel centres rather than texel boundaries
        uniforms.size.value.set(this.__composit_layer.renderTarget.width, this.__composit_layer.renderTarget.height);

        // setup camera
        const near = camera.near;
        const far = camera.far;

        uniforms.cameraNear.value = near;
        uniforms.cameraFar.value = far;

        uniforms.cameraInverseProjectionMatrix.value.copy(camera.projectionMatrixInverse);
        uniforms.cameraProjectionMatrix.value.copy(camera.projectionMatrix);


        // trigger uniform update
        material.uniformsNeedUpdate = true;

        // upscale runs at full resolution (the composit layer), reading the half-res AO target
        const upscale = this.__upscale_material;
        const upscaleUniforms = upscale.uniforms;
        const output = this.__composit_layer.renderTarget;

        upscaleUniforms.size.value.set(output.width, output.height);
        upscaleUniforms.cameraNear.value = near;
        upscaleUniforms.cameraFar.value = far;

        upscale.uniformsNeedUpdate = true;
    }

    __prepare_draw() {
        const engine = this.engine;
        const graphics = engine.graphics;
        const camera = graphics.camera;

        this.__setCamera(camera);

        this.__prepare_uniforms();

    }

    /**
     *
     * @private
     */
    __render() {
        const engine = this.engine;
        const graphics = engine.graphics;

        const renderer = graphics.renderer;

        this.__prepare_draw();

        const __old_state_rt = renderer.getRenderTarget();

        // render AO at half resolution
        renderer.setRenderTarget(this.__render_target);
        renderer.clearColor();

        renderScreenSpace(renderer, this.__material);

        // edge-aware upscale + denoise into the full-res composit layer (single non-separable pass;
        // the fullscreen draw writes every pixel, so no clear is required)
        renderer.setRenderTarget(this.__composit_layer.renderTarget);

        renderScreenSpace(renderer, this.__upscale_material);

        // restore rt
        renderer.setRenderTarget(__old_state_rt);
    }

    /**
     *
     * @param {boolean} value
     */
    setNormalBufferUsage(value) {
        if (value === this.__use_normals) {
            return;
        }

        this.__use_normals = value;

        const material = this.__material;
        const uniforms = material.uniforms;
        const defines = material.defines;

        const engine = this.engine;
        const graphics = engine.graphics;

        const buffer = graphics.frameBuffers.getById(StandardFrameBuffers.Normal);

        if (value) {
            uniforms.tNormal.value = buffer.getRenderTarget().texture;

            this.__add_used_frame_buffer(buffer);

            defines.NORMAL_TEXTURE = 1;
        } else {
            uniforms.tNormal.value = null;

            this.__remove_used_frame_buffer(buffer);
            defines.NORMAL_TEXTURE = 0;
        }

        // trigger re-compilation
        material.needsUpdate = true;
    }

    /**
     *
     * @param {FrameBuffer} buffer
     * @private
     */
    __add_used_frame_buffer(buffer) {
        if (this.__used_buffers.includes(buffer)) {
            // already used
            return;
        }

        buffer.referenceCount++;

        this.__used_buffers.push(buffer);
    }

    /**
     *
     * @param {FrameBuffer} buffer
     * @private
     */
    __remove_used_frame_buffer(buffer) {
        const i = this.__used_buffers.indexOf(buffer);
        if (i === -1) {
            // not used
            return;
        }

        buffer.referenceCount--;

        this.__used_buffers.splice(i, 1);
    }


    __update_render_target_size() {
        const engine = this.engine;
        const graphics = engine.graphics;

        // scale relative to the actual render resolution (device pixels), the same basis as the `size`
        // uniform and the full-res depth buffer the AO samples -- so the half-res grid stays a clean
        // fraction of the full-res grid at any devicePixelRatio / supersample factor
        const size = graphics.output_resolution;

        // clamp to 1: a 0-sized attachment trips GL_INVALID_FRAMEBUFFER_OPERATION on clear/draw
        const target_resolution_x = Math.max(1, Math.ceil(size.x * this.__resolution_scale));
        const target_resolution_y = Math.max(1, Math.ceil(size.y * this.__resolution_scale));

        this.__render_target.setSize(
            target_resolution_x,
            target_resolution_y
        );
    }

    async startup() {
        const engine = this.engine;
        const graphics = engine.graphics;

        // bind compositing layer
        this.__composit_layer = graphics.layerComposer.createLayer(
            {
                format: RedFormat,
                type: UnsignedByteType,
                internalFormat: 'R8'
            },
            BlendingType.Multiply
        );
        this.__composit_layer.name = 'Ambient Occlusion';

        // apply ambient occlusion on opaque pass, otherwise it would occlude transparent geometry
        this.__composit_layer.stage = CompositingStages.POST_OPAQUE;

        const material = this.__material;
        const uniforms = material.uniforms;

        const fb_color_depth = graphics.frameBuffers.getById(StandardFrameBuffers.ColorAndDepth);

        const depthTexture = fb_color_depth.getRenderTarget().depthTexture;
        uniforms.tDepth.value = depthTexture;

        if (this.__use_normals) {
            // force usage of normal buffer
            this.__use_normals = false;
            this.setNormalBufferUsage(true);
        }

        this.__composit_layer.on.preRender.add(this.__render, this);
        graphics.output_resolution.onChanged.add(this.__update_render_target_size, this);

        this.__update_render_target_size();

        // bind upscale inputs: half-res AO target + full-res depth
        const upscaleUniforms = this.__upscale_material.uniforms;
        upscaleUniforms.tAO.value = this.__render_target.texture;
        upscaleUniforms.tDepth.value = depthTexture;

        return super.startup();
    }

    async shutdown() {

        const engine = this.engine;
        const graphics = engine.graphics;

        // remove the layer
        graphics.layerComposer.remove(this.__composit_layer);

        this.__composit_layer.on.preRender.remove(this.__render, this);
        graphics.output_resolution.onChanged.remove(this.__update_render_target_size, this);

        // release memory
        this.__upscale_material.dispose();
        this.__hilbert_texture.dispose();
        this.__render_target.dispose();

        return super.shutdown();
    }

}
