import logger from '../../../virtual-background/logger';

import { ICompositeOptions, ICompositor } from './ICompositor';

const VERTEX_SHADER_SOURCE = `
    attribute vec2 a_position;
    attribute vec2 a_texCoord;
    varying vec2 v_texCoord;
    void main() {
        gl_Position = vec4(a_position, 0.0, 1.0);
        v_texCoord = a_texCoord;
    }
`;

// GLSL ES 1.0 — compatible with both WebGL 1 and WebGL 2.
// Temporal smoothing (EMA) is handled CPU-side by BackgroundFrameProcessor before upload,
// so the shader only reads the current mask — no ping-pong / prev-mask sampling needed.
const VB_FRAGMENT_SHADER_SOURCE = `
    precision mediump float;

    uniform sampler2D u_camera;
    uniform sampler2D u_background;
    uniform sampler2D u_mask;
    uniform float     u_edgeLow;
    uniform float     u_edgeHigh;
    uniform vec2      u_maskTexelSize;
    uniform vec2      u_cameraTexelSize;

    varying vec2 v_texCoord;

    // 3x3 Gaussian blur (sigma≈1) sampled in mask-texture space.
    // When u_maskTexelSize is (0,0) all taps collapse to tc — weights sum to 1.0, no blur.
    float sampleMaskBlurred(vec2 tc) {
        vec2 t = u_maskTexelSize;
        return
            texture2D(u_mask, tc + vec2(-t.x, -t.y)).r * 0.0625 +
            texture2D(u_mask, tc + vec2( 0.0, -t.y)).r * 0.125  +
            texture2D(u_mask, tc + vec2( t.x, -t.y)).r * 0.0625 +
            texture2D(u_mask, tc + vec2(-t.x,  0.0)).r * 0.125  +
            texture2D(u_mask, tc                    ).r * 0.25   +
            texture2D(u_mask, tc + vec2( t.x,  0.0)).r * 0.125  +
            texture2D(u_mask, tc + vec2(-t.x,  t.y)).r * 0.0625 +
            texture2D(u_mask, tc + vec2( 0.0,  t.y)).r * 0.125  +
            texture2D(u_mask, tc + vec2( t.x,  t.y)).r * 0.0625;
    }

    // Luminance gradient magnitude at this pixel using a simple cross-pattern sample.
    // Returns values in roughly [0, 0.5] — strong camera edges (e.g. shirt vs chair) hit 0.2+.
    float cameraEdgeStrength(vec2 tc) {
        vec2 t  = u_cameraTexelSize;
        vec3 lv = vec3(0.299, 0.587, 0.114);
        float lL = dot(texture2D(u_camera, tc - vec2(t.x, 0.0)).rgb, lv);
        float lR = dot(texture2D(u_camera, tc + vec2(t.x, 0.0)).rgb, lv);
        float lU = dot(texture2D(u_camera, tc - vec2(0.0, t.y)).rgb, lv);
        float lD = dot(texture2D(u_camera, tc + vec2(0.0, t.y)).rgb, lv);
        return length(vec2(lR - lL, lD - lU));
    }

    void main() {
        vec4 fg = texture2D(u_camera,     v_texCoord);
        vec4 bg = texture2D(u_background, v_texCoord);

        float blurred = sampleMaskBlurred(v_texCoord);
        float raw     = texture2D(u_mask, v_texCoord).r;

        // At strong camera color edges (e.g. dark chair vs bright shirt) snap toward the raw mask
        // so the person boundary aligns with the actual image edge. At smooth regions (hair, skin)
        // keep the blurred mask for natural feathering.
        float snapT = smoothstep(0.05, 0.20, cameraEdgeStrength(v_texCoord));
        float curAlpha = mix(blurred, raw, snapT);
        float alpha = smoothstep(u_edgeLow, u_edgeHigh, curAlpha);

        gl_FragColor = mix(bg, fg, alpha);
    }
`;

// Full-screen quad: two triangles covering clip space [-1,1].
const QUAD_VERTICES = new Float32Array([
    -1, -1, 0, 1,
    1, -1, 1, 1,
    -1, 1, 0, 0,
    1, 1, 1, 0
]);

function compileShader(gl: WebGLRenderingContext, type: number, source: string): WebGLShader {
    const shader = gl.createShader(type);

    if (!shader) {
        throw new Error('[VirtualBackground] WebGLCompositor: createShader returned null');
    }
    gl.shaderSource(shader, source);
    gl.compileShader(shader);
    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
        const info = gl.getShaderInfoLog(shader);

        gl.deleteShader(shader);
        throw new Error(`[VirtualBackground] WebGLCompositor: shader compile error — ${info}`);
    }

    return shader;
}

function createProgram(gl: WebGLRenderingContext, fragmentSource: string): WebGLProgram {
    const vert = compileShader(gl, gl.VERTEX_SHADER, VERTEX_SHADER_SOURCE);
    const frag = compileShader(gl, gl.FRAGMENT_SHADER, fragmentSource);
    const program = gl.createProgram();

    if (!program) {
        throw new Error('[VirtualBackground] WebGLCompositor: createProgram returned null');
    }
    gl.attachShader(program, vert);
    gl.attachShader(program, frag);
    gl.linkProgram(program);
    if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
        const info = gl.getProgramInfoLog(program);

        gl.deleteProgram(program);
        throw new Error(`[VirtualBackground] WebGLCompositor: program link error — ${info}`);
    }
    gl.deleteShader(vert);
    gl.deleteShader(frag);

    return program;
}

function createTexture(gl: WebGLRenderingContext): WebGLTexture {
    const tex = gl.createTexture();

    if (!tex) {
        throw new Error('[VirtualBackground] WebGLCompositor: createTexture returned null');
    }
    gl.bindTexture(gl.TEXTURE_2D, tex);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.LINEAR);

    return tex;
}

/**
 * WebGL-based compositor that blends camera, background, and segmentation mask
 * using a fragment shader with temporal smoothing and edge feathering (smoothstep).
 *
 * Falls back gracefully when context creation fails — callers should check
 * {@link WebGLCompositor#isAvailable} before use.
 */
export default class WebGLCompositor implements ICompositor {
    _canvas: HTMLCanvasElement;
    _gl: WebGLRenderingContext | null = null;

    // Bound event handlers — stored so they can be removed in dispose().
    _onContextLostBound: ((event: Event) => void) | null = null;
    _onContextRestoredBound: (() => void) | null = null;

    _program: WebGLProgram | null = null;
    _quadBuffer: WebGLBuffer | null = null;
    _texBackground: WebGLTexture | null = null;
    _texCamera: WebGLTexture | null = null;
    _texMask: WebGLTexture | null = null;

    // Attribute locations — resolved once after program link, reused every frame.
    // When sharing a GL context with TF.js the vertex attribute state (buffer
    // binding + vertexAttribPointer) is global and gets overwritten by TF.js
    // inference. Re-binding before drawArrays restores our state without the
    // overhead of re-querying locations each frame.
    _aPosition = -1;
    _aTexCoord = -1;

    // Uniform locations — VB program, resolved once after program link.
    _uBackground: WebGLUniformLocation | null = null;
    _uCamera: WebGLUniformLocation | null = null;
    _uCameraTexelSize: WebGLUniformLocation | null = null;
    _uEdgeHigh: WebGLUniformLocation | null = null;
    _uEdgeLow: WebGLUniformLocation | null = null;
    _uMask: WebGLUniformLocation | null = null;
    _uMaskTexelSize: WebGLUniformLocation | null = null;

    /**
     * Creates a WebGLCompositor targeting the given canvas.
     *
     * @param {HTMLCanvasElement} canvas - Output canvas element.
     */
    constructor(canvas: HTMLCanvasElement) {
        this._canvas = canvas;
        this._init();

        this._onContextLostBound = this._onContextLost.bind(this);
        this._onContextRestoredBound = this._onContextRestored.bind(this);
        canvas.addEventListener('webglcontextlost', this._onContextLostBound, false);
        canvas.addEventListener('webglcontextrestored', this._onContextRestoredBound, false);
    }

    /**
     * Returns true if a WebGL context is active and ready.
     *
     * @returns {boolean}
     */
    get isAvailable(): boolean {
        return this._gl !== null && !this._gl.isContextLost();
    }

    /**
     * Returns the canvas element that receives the composited output each frame.
     *
     * @returns {HTMLCanvasElement}
     */
    get outputCanvas(): HTMLCanvasElement {
        return this._canvas;
    }

    /**
     * Resizes the output canvas to the given dimensions.
     *
     * @param {number} width - New width in pixels.
     * @param {number} height - New height in pixels.
     * @returns {void}
     */
    resize(width: number, height: number): void {
        this._canvas.width = width;
        this._canvas.height = height;
    }

    /**
     * Composites one frame via the {@link ICompositor} interface.
     * Delegates to {@link compositeFromImageData}.
     *
     * @param {CanvasImageSource} camera - Live camera frame.
     * @param {CanvasImageSource} background - Pre-rendered background.
     * @param {ImageData} maskData - Raw segmentation mask.
     * @param {ICompositeOptions} options - Edge thresholds and blur radius.
     * @returns {void}
     */
    composite(
            camera: CanvasImageSource,
            background: CanvasImageSource,
            maskData: ImageData,
            options: ICompositeOptions): void {
        this.compositeFromImageData(
            camera as TexImageSource,
            background as HTMLCanvasElement,
            maskData,
            options.edgeLow,
            options.edgeHigh,
            options.maskBlurRadius
        );
    }

    /**
     * Composites using raw {@code ImageData} for the segmentation mask.
     *
     * The mask is uploaded via the {@code Uint8Array} form of {@code texImage2D} rather than
     * via the {@code CanvasImageSource} form. This is critical when the mask data comes from
     * {@code Mask.toImageData()} because the {@code CanvasImageSource} upload path reads from
     * the browser's internal canvas storage which uses <em>premultiplied alpha</em>: for this
     * model the R and A channels are both equal to the segmentation confidence, so the canvas
     * stores {@code R = confidence²} instead of {@code confidence}. The raw upload path
     * bypasses this distortion entirely.
     *
     * Temporal smoothing (EMA) is applied CPU-side by {@link BackgroundFrameProcessor} before
     * upload, so the shader only reads the current mask — no ping-pong required.
     *
     * @param {TexImageSource} cameraSource - Live camera frame (HTMLVideoElement, ImageBitmap, VideoFrame, etc.).
     * @param {HTMLCanvasElement} backgroundSource - Pre-rendered background canvas.
     * @param {ImageData} maskData - Current-frame segmentation mask (raw bytes).
     * @param {number} edgeLow - Smoothstep lower threshold.
     * @param {number} edgeHigh - Smoothstep upper threshold.
     * @param {number} maskBlurRadius - Gaussian blur radius in mask-texture pixels (0 = no blur).
     * Scaled from a target camera-pixel feathering width by the processor so the physical
     * blur width is consistent across tiers (TFLite uses a wider target to hide upscale artefacts).
     * @returns {void}
     */
    compositeFromImageData(
            cameraSource: TexImageSource,
            backgroundSource: HTMLCanvasElement,
            maskData: ImageData,
            edgeLow: number,
            edgeHigh: number,
            maskBlurRadius = 0
    ): void {
        const gl = this._gl;

        if (!gl || !this._program) {
            return;
        }

        gl.bindFramebuffer(gl.FRAMEBUFFER, null);
        gl.viewport(0, 0, this._canvas.width, this._canvas.height);
        gl.useProgram(this._program);

        // Camera texture
        gl.activeTexture(gl.TEXTURE0);
        gl.bindTexture(gl.TEXTURE_2D, this._texCamera);
        gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, cameraSource);
        gl.uniform1i(this._uCamera, 0);

        // Background texture
        gl.activeTexture(gl.TEXTURE1);
        gl.bindTexture(gl.TEXTURE_2D, this._texBackground);
        gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, backgroundSource);
        gl.uniform1i(this._uBackground, 1);

        // Current mask — raw Uint8Array upload bypasses browser premultiplied-alpha
        // conversion that would corrupt the R channel when uploading via canvas.
        gl.activeTexture(gl.TEXTURE2);
        gl.bindTexture(gl.TEXTURE_2D, this._texMask);
        gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, maskData.width, maskData.height,
            0, gl.RGBA, gl.UNSIGNED_BYTE, maskData.data);
        gl.uniform1i(this._uMask, 2);

        gl.uniform1f(this._uEdgeLow, edgeLow);
        gl.uniform1f(this._uEdgeHigh, edgeHigh);
        gl.uniform2f(this._uMaskTexelSize, maskBlurRadius / maskData.width, maskBlurRadius / maskData.height);
        gl.uniform2f(this._uCameraTexelSize, 1.0 / this._canvas.width, 1.0 / this._canvas.height);

        this._bindQuadState(gl);
        gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);
    }

    /**
     * Releases all WebGL resources.
     *
     * @returns {void}
     */
    dispose(): void {
        // Remove event listeners first so the canvas no longer holds a reference
        // to this compositor via the bound handler closures.
        if (this._onContextLostBound) {
            this._canvas.removeEventListener('webglcontextlost', this._onContextLostBound);
            this._onContextLostBound = null;
        }
        if (this._onContextRestoredBound) {
            this._canvas.removeEventListener('webglcontextrestored', this._onContextRestoredBound);
            this._onContextRestoredBound = null;
        }

        const gl = this._gl;

        if (!gl) {
            return;
        }
        gl.deleteTexture(this._texCamera);
        gl.deleteTexture(this._texBackground);
        gl.deleteTexture(this._texMask);
        gl.deleteBuffer(this._quadBuffer);
        gl.deleteProgram(this._program);

        this._texCamera = null;
        this._texBackground = null;
        this._texMask = null;
        this._quadBuffer = null;
        this._program = null;

        // Explicitly release the WebGL context so it no longer counts against the browser's
        // ~16 concurrent context limit. Without this the context stays alive until GC.
        gl.getExtension('WEBGL_lose_context')?.loseContext();
        this._gl = null;
    }

    /**
     * Initialises (or re-initialises) the WebGL context, shaders, and textures.
     *
     * @private
     * @returns {void}
     */
    _init(): void {
        const gl = (this._canvas.getContext('webgl2')
            ?? this._canvas.getContext('webgl')) as WebGLRenderingContext | null;

        if (!gl) {
            logger.warn('[VirtualBackground] WebGLCompositor: no WebGL context available — falling back to Canvas 2D');

            return;
        }
        this._gl = gl;

        try {
            this._program = createProgram(gl, VB_FRAGMENT_SHADER_SOURCE);
        } catch (err) {
            logger.error('[VirtualBackground] WebGLCompositor: VB shader build failed', err);
            this._gl = null;

            return;
        }

        // Vertex buffer — shared position + texCoord interleaved
        this._quadBuffer = gl.createBuffer();
        gl.bindBuffer(gl.ARRAY_BUFFER, this._quadBuffer);
        gl.bufferData(gl.ARRAY_BUFFER, QUAD_VERTICES, gl.STATIC_DRAW);

        this._aPosition = gl.getAttribLocation(this._program, 'a_position');
        this._aTexCoord = gl.getAttribLocation(this._program, 'a_texCoord');

        // Initial attribute setup — _bindQuadState() repeats this before every
        // draw call so TF.js cannot permanently corrupt our attribute state.
        this._bindQuadState(gl);

        // Resolve VB uniform locations.
        gl.useProgram(this._program);
        this._uCamera = gl.getUniformLocation(this._program, 'u_camera');
        this._uBackground = gl.getUniformLocation(this._program, 'u_background');
        this._uCameraTexelSize = gl.getUniformLocation(this._program, 'u_cameraTexelSize');
        this._uMask = gl.getUniformLocation(this._program, 'u_mask');
        this._uMaskTexelSize = gl.getUniformLocation(this._program, 'u_maskTexelSize');
        this._uEdgeLow = gl.getUniformLocation(this._program, 'u_edgeLow');
        this._uEdgeHigh = gl.getUniformLocation(this._program, 'u_edgeHigh');

        // Allocate textures
        this._texCamera = createTexture(gl);
        this._texBackground = createTexture(gl);
        this._texMask = createTexture(gl);

        logger.info('[VirtualBackground] WebGLCompositor: initialised');
    }

    /**
     * Handles WebGL context loss.
     *
     * @param {Event} event - The webglcontextlost event.
     * @private
     * @returns {void}
     */
    _onContextLost(event: Event): void {
        event.preventDefault();
        logger.warn('[VirtualBackground] WebGLCompositor: context lost');
        // The browser automatically destroys all GL objects on context loss.
        // Null the JS references so the GC can collect the wrapper objects.
        this._gl = null;
        this._program = null;
        this._quadBuffer = null;
        this._texCamera = null;
        this._texBackground = null;
        this._texMask = null;
    }

    /**
     * Handles WebGL context restoration.
     *
     * @private
     * @returns {void}
     */
    _onContextRestored(): void {
        logger.info('[VirtualBackground] WebGLCompositor: context restored — reinitialising');
        this._init();
    }

    /**
     * Re-binds the quad vertex buffer and re-specifies vertex attribute pointers.
     *
     * When the GL context is shared with TF.js, TF.js inference overwrites the
     * global vertex attribute state (bound ARRAY_BUFFER + vertexAttribPointer
     * descriptors). Calling this before every drawArrays() restores our state
     * without re-querying attribute locations.
     *
     * @private
     * @param {WebGLRenderingContext} gl - Active GL context.
     * @returns {void}
     */
    _bindQuadState(gl: WebGLRenderingContext): void {
        const stride = 4 * Float32Array.BYTES_PER_ELEMENT;

        gl.bindBuffer(gl.ARRAY_BUFFER, this._quadBuffer);
        gl.enableVertexAttribArray(this._aPosition);
        gl.vertexAttribPointer(this._aPosition, 2, gl.FLOAT, false, stride, 0);
        gl.enableVertexAttribArray(this._aTexCoord);
        gl.vertexAttribPointer(this._aTexCoord, 2, gl.FLOAT, false, stride, 2 * Float32Array.BYTES_PER_ELEMENT);
    }
}