web/lib/gpu-augmentation.ts

// GPU-accelerated augmentation using WebGL
// Applies brightness, contrast, and noise operations on the GPU

// Check if WebGL is available
let webglAvailable: boolean | null = null
let glContext: WebGLRenderingContext | null = null
let glCanvas: HTMLCanvasElement | null = null

export function isWebGLAvailable(): boolean {
    if (webglAvailable !== null) return webglAvailable

    try {
        const canvas = document.createElement('canvas')
        const gl = canvas.getContext('webgl') || canvas.getContext('experimental-webgl')
        webglAvailable = gl !== null
        return webglAvailable
    } catch {
        webglAvailable = false
        return false
    }
}

// Shader source code
const vertexShaderSource = `
    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;
    }
`

const fragmentShaderSource = `
    precision mediump float;
    
    uniform sampler2D u_image;
    uniform float u_brightness;
    uniform float u_contrast;
    uniform float u_noiseAmount;
    uniform float u_seed;
    
    varying vec2 v_texCoord;
    
    // Simple pseudo-random function
    float random(vec2 co) {
        return fract(sin(dot(co.xy + u_seed, vec2(12.9898, 78.233))) * 43758.5453);
    }
    
    void main() {
        vec4 color = texture2D(u_image, v_texCoord);
        
        // Apply brightness (additive)
        color.rgb += u_brightness;
        
        // Apply contrast (multiplicative from midpoint)
        color.rgb = (color.rgb - 0.5) * u_contrast + 0.5;
        
        // Apply noise
        if (u_noiseAmount > 0.0) {
            float noise = (random(v_texCoord) - 0.5) * u_noiseAmount;
            color.rgb += noise;
        }
        
        // Clamp values
        color.rgb = clamp(color.rgb, 0.0, 1.0);
        
        gl_FragColor = color;
    }
`

// Compile shader
function compileShader(gl: WebGLRenderingContext, source: string, type: number): WebGLShader | null {
    const shader = gl.createShader(type)
    if (!shader) return null

    gl.shaderSource(shader, source)
    gl.compileShader(shader)

    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
        console.error('Shader compile error:', gl.getShaderInfoLog(shader))
        gl.deleteShader(shader)
        return null
    }

    return shader
}

// Create shader program
function createProgram(gl: WebGLRenderingContext): WebGLProgram | null {
    const vertexShader = compileShader(gl, vertexShaderSource, gl.VERTEX_SHADER)
    const fragmentShader = compileShader(gl, fragmentShaderSource, gl.FRAGMENT_SHADER)

    if (!vertexShader || !fragmentShader) return null

    const program = gl.createProgram()
    if (!program) return null

    gl.attachShader(program, vertexShader)
    gl.attachShader(program, fragmentShader)
    gl.linkProgram(program)

    if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
        console.error('Program link error:', gl.getProgramInfoLog(program))
        gl.deleteProgram(program)
        return null
    }

    return program
}

// Initialize WebGL context and program
let program: WebGLProgram | null = null
let positionBuffer: WebGLBuffer | null = null
let texCoordBuffer: WebGLBuffer | null = null

function initWebGL(width: number, height: number): boolean {
    if (glCanvas && glContext && program) {
        // Resize if needed
        if (glCanvas.width !== width || glCanvas.height !== height) {
            glCanvas.width = width
            glCanvas.height = height
            glContext.viewport(0, 0, width, height)
        }
        return true
    }

    try {
        glCanvas = document.createElement('canvas')
        glCanvas.width = width
        glCanvas.height = height

        glContext = glCanvas.getContext('webgl', { preserveDrawingBuffer: true }) as WebGLRenderingContext
        if (!glContext) {
            console.warn('WebGL not available')
            return false
        }

        program = createProgram(glContext)
        if (!program) {
            console.warn('Failed to create WebGL program')
            return false
        }

        // Create buffers
        const gl = glContext

        // Position buffer (full screen quad)
        positionBuffer = gl.createBuffer()
        gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer)
        gl.bufferData(gl.ARRAY_BUFFER, new Float32Array([
            -1, -1, 1, -1, -1, 1,
            -1, 1, 1, -1, 1, 1
        ]), gl.STATIC_DRAW)

        // Texture coordinate buffer
        texCoordBuffer = gl.createBuffer()
        gl.bindBuffer(gl.ARRAY_BUFFER, texCoordBuffer)
        gl.bufferData(gl.ARRAY_BUFFER, new Float32Array([
            0, 1, 1, 1, 0, 0,
            0, 0, 1, 1, 1, 0
        ]), gl.STATIC_DRAW)

        return true
    } catch (err) {
        console.warn('WebGL initialization failed:', err)
        return false
    }
}

export interface GPUAugmentOptions {
    brightness?: number  // -1 to 1 (0 = no change)
    contrast?: number    // 0 to 2 (1 = no change)
    noiseAmount?: number // 0 to 1
    seed?: number
}

/**
 * Apply GPU-accelerated augmentation to an image
 * Returns the augmented canvas, or null if WebGL is not available
 */
export function applyGPUAugmentation(
    sourceCanvas: HTMLCanvasElement,
    options: GPUAugmentOptions
): HTMLCanvasElement | null {
    const { brightness = 0, contrast = 1, noiseAmount = 0, seed = Math.random() } = options

    // Skip if no augmentation needed
    if (brightness === 0 && contrast === 1 && noiseAmount === 0) {
        return sourceCanvas
    }

    if (!initWebGL(sourceCanvas.width, sourceCanvas.height)) {
        return null // Fallback to CPU
    }

    const gl = glContext!

    try {
        gl.useProgram(program!)

        // Create texture from source canvas
        const texture = gl.createTexture()
        gl.bindTexture(gl.TEXTURE_2D, texture)
        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)
        gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, sourceCanvas)

        // Set up attributes
        const positionLocation = gl.getAttribLocation(program!, 'a_position')
        const texCoordLocation = gl.getAttribLocation(program!, 'a_texCoord')

        gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer!)
        gl.enableVertexAttribArray(positionLocation)
        gl.vertexAttribPointer(positionLocation, 2, gl.FLOAT, false, 0, 0)

        gl.bindBuffer(gl.ARRAY_BUFFER, texCoordBuffer!)
        gl.enableVertexAttribArray(texCoordLocation)
        gl.vertexAttribPointer(texCoordLocation, 2, gl.FLOAT, false, 0, 0)

        // Set uniforms
        gl.uniform1f(gl.getUniformLocation(program!, 'u_brightness'), brightness)
        gl.uniform1f(gl.getUniformLocation(program!, 'u_contrast'), contrast)
        gl.uniform1f(gl.getUniformLocation(program!, 'u_noiseAmount'), noiseAmount)
        gl.uniform1f(gl.getUniformLocation(program!, 'u_seed'), seed)

        // Draw
        gl.drawArrays(gl.TRIANGLES, 0, 6)

        // Clean up texture
        gl.deleteTexture(texture)

        return glCanvas!
    } catch (err) {
        console.warn('GPU augmentation failed:', err)
        return null
    }
}

/**
 * Apply GPU augmentation and return the result as ImageData
 */
export function applyGPUAugmentationToImageData(
    imageData: ImageData,
    options: GPUAugmentOptions
): ImageData | null {
    // Create temporary canvas with the image data
    const tempCanvas = document.createElement('canvas')
    tempCanvas.width = imageData.width
    tempCanvas.height = imageData.height
    const ctx = tempCanvas.getContext('2d')
    if (!ctx) return null

    ctx.putImageData(imageData, 0, 0)

    const resultCanvas = applyGPUAugmentation(tempCanvas, options)
    if (!resultCanvas) return null

    const resultCtx = resultCanvas.getContext('2d') || resultCanvas.getContext('webgl')

    // For WebGL canvas, we need to read pixels differently
    if (glContext) {
        const pixels = new Uint8Array(imageData.width * imageData.height * 4)
        glContext.readPixels(0, 0, imageData.width, imageData.height, glContext.RGBA, glContext.UNSIGNED_BYTE, pixels)

        // Flip Y axis (WebGL has origin at bottom-left)
        const flipped = new Uint8ClampedArray(pixels.length)
        const rowSize = imageData.width * 4
        for (let y = 0; y < imageData.height; y++) {
            const srcRow = (imageData.height - 1 - y) * rowSize
            const dstRow = y * rowSize
            flipped.set(pixels.subarray(srcRow, srcRow + rowSize), dstRow)
        }

        return new ImageData(flipped, imageData.width, imageData.height)
    }

    return null
}

/**
 * Clean up WebGL resources
 */
export function cleanupGPU() {
    if (glContext && program) {
        glContext.deleteProgram(program)
        if (positionBuffer) glContext.deleteBuffer(positionBuffer)
        if (texCoordBuffer) glContext.deleteBuffer(texCoordBuffer)
    }
    program = null
    positionBuffer = null
    texCoordBuffer = null
    glContext = null
    glCanvas = null
}