import { Camera, DynamicTexture, Effect, PostProcess, Texture, Vector3 } from '../loader/babylonjs-import';
import { ssaocalculatorShader } from '../loader/CustomBabylonMaterials/shaders/SSAOCalculatorShader';
import { EXPENSIVE_NORMALS } from '../loader/ShaderConstants';
import { ShortFrustumDepthRenderer } from '../loader/ShortFrustumDepthRenderer';
import { TwinfinityPostProcess } from './TwinfinityPostProcess';

const SHADER_NAME_PREFIX = 'custom_ssao';
const DEPTH_SAMPLER = 'depthSampler';
const NOISE_SAMPLER = 'noiseSampler';

const DEPTH_DIMENSIONS = 'depthDimensions';
const INVERTED_PROJECTION_MATRIX = 'invertedProjectionMatrix';
const NEAR_FAR = 'nearFar';
const SAMPLE_KERNEL_SIZE = 'sampleKernelSize';
const SAMPLE_KERNEL = 'sampleKernel';
const STRENGTH = 'strength';
const NOISE_SCALE = 'noiseScale';
const UV_SCALE = 'scale';
const RADIUS = 'radius';
const PROJECTION_MATRIX = 'projectionMatrix';

/**
 * Calculates SSAO using the depth buffer.
 * Pretty much John Chapmans SSAO tutorial: http://john-chapman-graphics.blogspot.com/2013/01/ssao-tutorial.html but there is no normal buffer so we have to recreate the normals from the depth buffer instead
 */
export class SSAOCalculator extends TwinfinityPostProcess {
    private _shortFrustumDepthRenderer: ShortFrustumDepthRenderer;

    private _kernel: Vector3[] = [];
    private _noiseTexture: DynamicTexture;
    public strength: number;
    public radius: number;
    private _samples: number;
    private readonly _renderScale: number;

    public set samples(samples: number) {
        this._samples = samples;
        this.createKernel();

        if (this._postProcess) {
            this._postProcess.updateEffect(
                this.getDefines(),
                this.getUniforms(),
                this.getSamplers(),
                undefined,
                () => {},
                undefined
            );
        } else {
            throw new Error('Updating samples of uinitialized ssao calculator');
        }
    }

    public get samples(): number {
        return this._samples;
    }

    private getDefines(): string {
        const defines: string[] = [];
        defines.push(`#define ${SAMPLE_KERNEL_SIZE} ${this.samples}`);
        defines.push(`#define ${EXPENSIVE_NORMALS}`);
        return defines.join('\n');
    }

    private getUniforms(): string[] {
        const uniforms = [
            UV_SCALE,
            STRENGTH,
            RADIUS,
            NOISE_SCALE,
            DEPTH_DIMENSIONS,
            NEAR_FAR,
            PROJECTION_MATRIX,
            INVERTED_PROJECTION_MATRIX
        ];

        for (let i = 0; i < this._kernel.length; i++) {
            uniforms.push(SAMPLE_KERNEL + '[' + i + ']');
        }

        return uniforms;
    }

    private getSamplers(): string[] {
        return [DEPTH_SAMPLER, NOISE_SAMPLER];
    }

    protected initialize(camera: Camera): PostProcess {
        const engine = camera.getEngine();

        const scene = camera.getScene();

        const noiseSize = 64;
        this._noiseTexture = new DynamicTexture(
            'SSAOnoiseTexture',
            noiseSize,
            scene,
            false,
            Texture.NEAREST_SAMPLINGMODE
        );
        this._noiseTexture.wrapR = 1;
        this._noiseTexture.wrapU = 1;
        this._noiseTexture.wrapV = 1;

        const noiseTextureContext = this._noiseTexture.getContext();

        for (let i = 0; i < noiseSize; i++) {
            for (let j = 0; j < noiseSize; j++) {
                noiseTextureContext.fillStyle = `rgba(
                    ${Math.floor(Math.random() * 255.0)},
                    ${Math.floor(Math.random() * 255.0)},
                    0)`;
                noiseTextureContext.fillRect(i, j, 1, 1);
            }
        }
        this._noiseTexture.update();

        if (this._postProcess === undefined) {
            Effect.ShadersStore[SHADER_NAME_PREFIX + 'FragmentShader'] = ssaocalculatorShader.fragment();
            const newPostProcess = new PostProcess(
                'SSAO ',
                SHADER_NAME_PREFIX,
                this.getUniforms(),
                this.getSamplers(),
                {
                    // Do not include scale here, because that would mean that the postprocess before this one (this blit) would be resized, because Babylonjs has really weird rules for this
                    width: engine.getRenderWidth(),
                    height: engine.getRenderHeight()
                },
                null,
                undefined,
                engine,
                false,
                this.getDefines()
            );

            const viewer = camera.getScene().twinfinity.viewer;
            if (viewer) {
                const nonLinearDepth = true;
                this._shortFrustumDepthRenderer = viewer.enableDepthRenderer(nonLinearDepth);
                const shortFrustumProjectionMatrixState =
                    this._shortFrustumDepthRenderer.shortFrustumProjectionMatrixState;

                newPostProcess.onApplyObservable.add((effect) => {
                    effect.setTexture(NOISE_SAMPLER, this._noiseTexture);
                    effect.setTexture(DEPTH_SAMPLER, this._shortFrustumDepthRenderer.depthRenderer.getDepthMap());
                    effect.setFloat(STRENGTH, this.strength);
                    effect.setFloat(RADIUS, this.radius);

                    const renderWidth = engine.getRenderWidth();
                    const renderHeight = engine.getRenderHeight();
                    const depthDimensionsX = 1.0 / (renderWidth * this._renderScale);
                    const depthDimensionsY = 1.0 / (renderHeight * this._renderScale);
                    effect.setFloat2(DEPTH_DIMENSIONS, depthDimensionsX, depthDimensionsY);
                    effect.setFloat2(
                        NOISE_SCALE,
                        (renderWidth * this._renderScale) / noiseSize,
                        (renderHeight * this._renderScale) / noiseSize
                    );
                    effect.setFloat2(UV_SCALE, 1.0, 1.0);
                    effect.setFloat2(
                        NEAR_FAR,
                        shortFrustumProjectionMatrixState.depthRenderCameraMinZ,
                        shortFrustumProjectionMatrixState.depthRenderCameraMaxZ
                    );

                    effect.setMatrix(
                        INVERTED_PROJECTION_MATRIX,
                        shortFrustumProjectionMatrixState.getInvertedProjectionMatrix()
                    );
                    effect.setMatrix(PROJECTION_MATRIX, shortFrustumProjectionMatrixState.getProjectionMatrix());

                    for (let i = 0; i < this._kernel.length; i++) {
                        const kernelSample = this._kernel[i];
                        const sampleKernelUniformName = `${SAMPLE_KERNEL}[${i}]`;
                        effect.setFloat3(sampleKernelUniformName, kernelSample.x, kernelSample.y, kernelSample.z);
                    }
                });
            } else {
                throw new Error('No viewer created?');
            }

            this._postProcess = newPostProcess;

            return newPostProcess;
        } else return this._postProcess;
    }

    private createKernel(): void {
        const lerp = function (a: number, b: number, n: number): number {
            return (1 - n) * a + n * b;
        };

        this._kernel = [];

        for (let i = 0; i < this.samples; i++) {
            const kernelSample = new Vector3(Math.random() * 2.0 - 1.0, Math.random() * 2.0 - 1.0, Math.random());
            kernelSample.normalize();

            let scale = (i % this.samples) / this.samples;
            scale = lerp(0.1, 1.0, scale * scale);

            kernelSample.scaleInPlace(scale);

            this._kernel.push(kernelSample);
        }
    }

    constructor(samplesCount: number, strength: number, radius: number, renderScale: number) {
        super();
        this._samples = samplesCount;
        this.strength = strength;
        this.radius = radius;
        this._renderScale = renderScale;
        this.createKernel();
    }

    detach(camera: Camera): boolean {
        const detached = super.detach(camera);

        if (detached) {
            this._shortFrustumDepthRenderer.disable();
        }

        return detached;
    }
}
