import { AbstractPostEffect, PostEffectLayer } from './posteffect.js';
import { linearToGamma } from '../shaders/misc.js';
import { fetchSampler } from '../utility/misc.js';
import { Vector2 } from '@zephyr3d/base';
import '../render/primitive.js';
import { BoxShape } from '../shapes/box.js';
import '../shapes/capsule.js';
import '../shapes/cylinder.js';
import '../shapes/torus.js';
import '../shapes/plane.js';
import '../shapes/sphere.js';
import '../shapes/tetrahedron.js';
import { temporalResolve } from '../shaders/temporal.js';
import { RGHistoryResources } from '../render/rendergraph/history_resources.js';

/** @internal */ class TAA extends AbstractPostEffect {
    static _resolveProgram = [];
    static _skyMotionVectorProgram = null;
    static _box;
    static _texSize = new Vector2();
    _bindGroup;
    _skyMotionVectorBindGroup;
    constructor(){
        super();
        this._bindGroup = null;
        this._layer = PostEffectLayer.end;
        this._skyMotionVectorBindGroup = null;
    }
    renderSkyMotionVectors(ctx) {
        const fb = ctx.device.pool.fetchTemporalFramebuffer(false, 0, 0, ctx.motionVectorTexture, ctx.depthTexture);
        const program = TAA._getSkyMotionVectorProgram(ctx);
        if (!this._skyMotionVectorBindGroup) {
            this._skyMotionVectorBindGroup = ctx.device.createBindGroup(program.bindGroupLayouts[0]);
        }
        const box = TAA._getBox(ctx);
        this._skyMotionVectorBindGroup.setValue('VPMatrix', ctx.camera.viewProjectionMatrix);
        this._skyMotionVectorBindGroup.setValue('prevVPMatrix', ctx.camera.prevVPMatrix);
        this._skyMotionVectorBindGroup.setValue('cameraPos', ctx.camera.getWorldPosition());
        this._skyMotionVectorBindGroup.setValue('prevCameraPos', ctx.camera.prevPosition);
        ctx.device.pushDeviceStates();
        ctx.device.setProgram(program);
        ctx.device.setBindGroup(0, this._skyMotionVectorBindGroup);
        ctx.device.setRenderStates(AbstractPostEffect.getDefaultRenderState(ctx, 'le'));
        ctx.device.setFramebuffer(fb);
        box.draw();
        ctx.device.popDeviceStates();
        ctx.device.pool.releaseFrameBuffer(fb);
    }
    apply(ctx, inputColorTexture, sceneDepthTexture, srgbOutput) {
        const historyMgr = ctx.camera.getHistoryResourceManager();
        const colorDesc = {
            format: inputColorTexture.format,
            sizeMode: 'absolute',
            width: inputColorTexture.width,
            height: inputColorTexture.height
        };
        const colorSize = {
            width: inputColorTexture.width,
            height: inputColorTexture.height
        };
        const motionVectorDesc = {
            format: ctx.motionVectorTexture.format,
            sizeMode: 'absolute',
            width: ctx.motionVectorTexture.width,
            height: ctx.motionVectorTexture.height
        };
        const motionVectorSize = {
            width: ctx.motionVectorTexture.width,
            height: ctx.motionVectorTexture.height
        };
        const hasHistory = historyMgr.isCompatible(RGHistoryResources.TAA_COLOR, colorDesc, colorSize) && historyMgr.isCompatible(RGHistoryResources.TAA_MOTION_VECTOR, motionVectorDesc, motionVectorSize);
        if (!hasHistory) {
            // First frame: just pass through
            this.passThrough(ctx, inputColorTexture, srgbOutput);
            // Queue history resources only after retaining references from the pool.
            const currentColorTex = ctx.device.getFramebuffer().getColorAttachments()[0];
            ctx.device.pool.retainTexture(currentColorTex);
            historyMgr.queueCommit(RGHistoryResources.TAA_COLOR, colorDesc, colorSize, currentColorTex);
            ctx.device.pool.retainTexture(ctx.motionVectorTexture);
            historyMgr.queueCommit(RGHistoryResources.TAA_MOTION_VECTOR, motionVectorDesc, motionVectorSize, ctx.motionVectorTexture);
            return;
        }
        // Get history textures
        const prevColorTex = historyMgr.getPrevious(RGHistoryResources.TAA_COLOR);
        const prevMotionVectorTex = historyMgr.getPrevious(RGHistoryResources.TAA_MOTION_VECTOR);
        // Render TAA
        let program = TAA._resolveProgram[ctx.camera.TAADebug];
        if (!program) {
            program = TAA._getResolveProgram(ctx, ctx.camera.TAADebug);
            TAA._resolveProgram[ctx.camera.TAADebug] = program;
        }
        if (!this._bindGroup) {
            this._bindGroup = ctx.device.createBindGroup(program.bindGroupLayouts[0]);
        }
        this._bindGroup.setTexture('historyColorTex', prevColorTex, fetchSampler('clamp_linear_nomip'));
        this._bindGroup.setTexture('currentColorTex', inputColorTexture, fetchSampler('clamp_nearest_nomip'));
        this._bindGroup.setTexture('currentDepthTex', sceneDepthTexture, fetchSampler('clamp_nearest_nomip'));
        this._bindGroup.setTexture('motionVector', ctx.motionVectorTexture, fetchSampler('clamp_nearest_nomip'));
        this._bindGroup.setTexture('prevMotionVector', prevMotionVectorTex, fetchSampler('clamp_nearest_nomip'));
        this._bindGroup.setValue('flip', this.needFlip(ctx.device) ? 1 : 0);
        this._bindGroup.setValue('srgbOut', srgbOutput ? 1 : 0);
        TAA._texSize.setXY(sceneDepthTexture.width, sceneDepthTexture.height);
        this._bindGroup.setValue('texSize', TAA._texSize);
        ctx.device.setProgram(program);
        ctx.device.setBindGroup(0, this._bindGroup);
        this.drawFullscreenQuad();
        // Update history with current framebuffer output (zero-copy)
        const currentColorTex = ctx.device.getFramebuffer().getColorAttachments()[0];
        ctx.device.pool.retainTexture(currentColorTex);
        historyMgr.queueCommit(RGHistoryResources.TAA_COLOR, colorDesc, colorSize, currentColorTex);
        ctx.device.pool.retainTexture(ctx.motionVectorTexture);
        historyMgr.queueCommit(RGHistoryResources.TAA_MOTION_VECTOR, motionVectorDesc, motionVectorSize, ctx.motionVectorTexture);
    // History commits happen after the render graph completes successfully.
    }
    requireLinearDepthTexture(_ctx) {
        return true;
    }
    requireDepthAttachment(_ctx) {
        return true;
    }
    static _getSkyMotionVectorProgram(ctx) {
        if (!this._skyMotionVectorProgram) {
            this._skyMotionVectorProgram = ctx.device.buildRenderProgram({
                vertex (pb) {
                    this.$inputs.pos = pb.vec3().attrib('position');
                    this.VPMatrix = pb.mat4().uniform(0);
                    this.prevVPMatrix = pb.mat4().uniform(0);
                    this.cameraPos = pb.vec3().uniform(0);
                    this.prevCameraPos = pb.vec3().uniform(0);
                    pb.main(function() {
                        this.$l.worldPos = pb.add(this.$inputs.pos, this.cameraPos);
                        this.$l.prevWorldPos = pb.add(this.$inputs.pos, this.prevCameraPos);
                        this.$l.clipPos = pb.mul(this.VPMatrix, pb.vec4(this.worldPos, 1));
                        this.$l.prevClipPos = pb.mul(this.prevVPMatrix, pb.vec4(this.prevWorldPos, 1));
                        this.clipPos.z = this.clipPos.w;
                        this.$builtins.position = this.clipPos;
                        this.$outputs.currentPos = this.clipPos;
                        this.$outputs.prevPos = this.prevClipPos;
                    });
                },
                fragment (pb) {
                    this.$outputs.color = pb.vec4();
                    pb.main(function() {
                        this.$l.motionVector = pb.mul(pb.sub(pb.div(this.$inputs.currentPos.xy, this.$inputs.currentPos.w), pb.div(this.$inputs.prevPos.xy, this.$inputs.prevPos.w)), 0.5);
                        this.$outputs.color = pb.vec4(this.motionVector, 0, 1);
                    });
                }
            });
            this._skyMotionVectorProgram.name = '@TAA_SkyMotionVector';
        }
        return this._skyMotionVectorProgram;
    }
    static _getBox(_ctx) {
        if (!this._box) {
            this._box = new BoxShape({
                size: 2,
                needNormal: false,
                needUV: false
            });
        }
        return this._box;
    }
    static _getResolveProgram(ctx, debug) {
        const program = ctx.device.buildRenderProgram({
            vertex (pb) {
                this.flip = pb.int().uniform(0);
                this.$inputs.pos = pb.vec2().attrib('position');
                this.$outputs.uv = pb.vec2();
                pb.main(function() {
                    this.$builtins.position = pb.vec4(this.$inputs.pos, 1, 1);
                    this.$outputs.uv = pb.add(pb.mul(this.$inputs.pos.xy, 0.5), pb.vec2(0.5));
                    this.$if(pb.notEqual(this.flip, 0), function() {
                        this.$builtins.position.y = pb.neg(this.$builtins.position.y);
                    });
                });
            },
            fragment (pb) {
                this.historyColorTex = pb.tex2D().uniform(0);
                this.currentColorTex = pb.tex2D().uniform(0);
                this.currentDepthTex = pb.tex2D().uniform(0);
                this.motionVector = pb.tex2D().uniform(0);
                this.prevMotionVector = pb.tex2D().uniform(0);
                this.texSize = pb.vec2().uniform(0);
                this.srgbOut = pb.int().uniform(0);
                this.$outputs.outColor = pb.vec4();
                pb.main(function() {
                    this.$l.screenUV = pb.div(pb.vec2(this.$builtins.fragCoord.xy), this.texSize);
                    this.$l.resolvedColor = temporalResolve(this, this.currentColorTex, this.historyColorTex, this.currentDepthTex, this.motionVector, this.prevMotionVector, this.screenUV, this.texSize, debug);
                    this.$if(pb.equal(this.srgbOut, 0), function() {
                        this.$outputs.outColor = pb.vec4(this.resolvedColor, 1);
                    }).$else(function() {
                        this.$outputs.outColor = pb.vec4(linearToGamma(this, this.resolvedColor), 1);
                    });
                });
            }
        });
        program.name = '@TAA_Resolve';
        return program;
    }
}

export { TAA };
//# sourceMappingURL=taa.js.map
