# ndwasm_image.js

    import { NDArray, DTYPE_MAP } from './ndarray_core.js';
    import NDWasm from './ndwasm.js';

    /**
     * @namespace NDWasmImage
     * @description Provides WebAssembly-powered functions for image processing.
     * This module allows for efficient conversion between image binary data and NDArrays.
     */
    export const NDWasmImage = {
        /**
         * Decodes a binary image into an NDArray.
         * Supports common formats like PNG, JPEG, GIF, and WebP.
         * The resulting NDArray will have a shape of [height, width, 4] and a 'uint8c' dtype,
         * representing RGBA channels. This provides a consistent starting point for image manipulation.
         *
         * @memberof NDWasmImage
         * @param {Uint8Array} imageBytes - The raw binary data of the image file.
         * @returns {NDArray|null} A 3D NDArray representing the image, or null if decoding fails.
         * @example
         * const imageBlob = await fetch('./my-image.png').then(res => res.blob());
         * const imageBytes = new Uint8Array(await imageBlob.arrayBuffer());
         * const imageArray = NDWasmImage.decode(imageBytes);
         * // imageArray.shape is [height, width, 4]
         */
        decode(imageBytes) {
            if (!NDWasm.runtime || !NDWasm.runtime.isLoaded) {
                throw new Error("NDWasm runtime is not initialized. Call NDWasm.init() first.");
            }
            if (!(imageBytes instanceof Uint8Array)) {
                throw new Error("Input must be a Uint8Array.");
            }

            const exports = NDWasm.runtime.exports;
            let imagePtr = 0;
            let resultPtr = 0;
            let pixelDataPtr = 0;

            try {
                // 1. Allocate memory for the input image and copy it to WASM.
                imagePtr = exports.malloc(imageBytes.length);
                if (!imagePtr) throw new Error("WASM malloc failed for image bytes.");
                new Uint8Array(exports.mem.buffer, imagePtr, imageBytes.length).set(imageBytes);

                // 2. Call the WASM decoding function.
                resultPtr = exports.decode_image(imagePtr, imageBytes.length);
                if (!resultPtr) {
                    console.error("Failed to decode image in WASM. The image format may be invalid or unsupported.");
                    return null;
                }

                // 3. Read the result structure [pixel_ptr, pixel_size, width, height] from WASM memory.
                const resultView = new DataView(exports.mem.buffer, resultPtr, 16);
                pixelDataPtr = resultView.getUint32(0, true);
                const pixelDataSize = resultView.getUint32(4, true);
                const width = resultView.getUint32(8, true);
                const height = resultView.getUint32(12, true);

                if (!pixelDataPtr) {
                    console.error("WASM decode_image returned a null pixel data pointer.");
                    return null;
                }

                // 4. Create a copy of the pixel data from WASM memory.
                const pixelData = new Uint8ClampedArray(exports.mem.buffer, pixelDataPtr, pixelDataSize).slice();

                // 5. Create the final NDArray.
                return new NDArray(pixelData, { shape: [height, width, 4], dtype: 'uint8c' });

            } finally {
                // 6. Free all allocated WASM memory.
                if (imagePtr) exports.free(imagePtr, imageBytes.length);
                if (resultPtr) exports.free(resultPtr, 16);
                if (pixelDataPtr) {
                    const resultView = new DataView(exports.mem.buffer, resultPtr, 16);
                    const pixelDataSize = resultView.getUint32(4, true);
                    exports.free(pixelDataPtr, pixelDataSize);
                }
            }
        },

        /**
         * Encodes an NDArray into a binary image format (PNG or JPEG).
         *
         * @memberof NDWasmImage
         * @param {NDArray} ndarray - The input array.
         *   Supported dtypes: 'uint8', 'uint8c', 'float32', 'float64'. Float values should be in the range [0, 1].
         *   Supported shapes: [h, w] (grayscale), [h, w, 1] (grayscale), [h, w, 3] (RGB), or [h, w, 4] (RGBA).
         * @param {object} [options={}] - Encoding options.
         * @param {string} [options.format='png'] - The target format: 'png' or 'jpeg'. It is recommended to use the `encodePng` or `encodeJpeg` helpers instead.
         * @param {number} [options.quality=90] - The quality for JPEG encoding (1-100). Ignored for PNG.
         * @returns {Uint8Array|null} A Uint8Array containing the binary data of the encoded image, or null on failure.
         * @see {@link NDWasmImage.encodePng}
         * @see {@link NDWasmImage.encodeJpeg}
         * @example
         * // Encode a 3-channel float array to a high-quality JPEG using the main function
         * const floatArr = NDArray.random([100, 150, 3]);
         * const jpegBytes = NDWasmImage.encode(floatArr, { format: 'jpeg', quality: 95 });
         */
        encode(ndarray, { format = 'png', quality = 90 } = {}) {
            if (!NDWasm.runtime || !NDWasm.runtime.isLoaded) {
                throw new Error("NDWasm runtime is not initialized. Call NDWasm.init() first.");
            }
            if (!(ndarray instanceof NDArray)) {
                throw new Error("Input must be an NDArray.");
            }

            // --- 1. Validate and Pre-process NDArray ---

            // Validate dimensions and channels
            let channels;
            if (ndarray.ndim === 2) {
                channels = 1;
            } else if (ndarray.ndim === 3) {
                channels = ndarray.shape[2];
                if (channels !== 1 && channels !== 3 && channels !== 4) {
                    throw new Error(`Unsupported channel count for 3D array: ${channels}. Must be 1, 3, or 4.`);
                }
            } else {
                throw new Error(`Unsupported array dimensions: ${ndarray.ndim}. Must be 2 or 3.`);
            }

            let arrayToEncode = ndarray;

            // Handle float dtypes by scaling to uint8
            if (ndarray.dtype === 'float32' || ndarray.dtype === 'float64') {
                const scaledData = new Uint8Array(ndarray.size);
                // This is a simple but slow way to iterate. A faster, vectorized `mul` and `astype` would be an improvement.
                let i = 0;
                ndarray.iterate(value => {
                    scaledData[i++] = Math.max(0, Math.min(255, value * 255));
                });
                arrayToEncode = new NDArray(scaledData, { shape: ndarray.shape, dtype: 'uint8' });
            }

            // Ensure data is contiguous C-style for easy passing to WASM
            let sourceData = arrayToEncode.data;
            if (!arrayToEncode.isContiguous) {
                const flatData = new (DTYPE_MAP[arrayToEncode.dtype])(arrayToEncode.size);
                let i = 0;
                arrayToEncode.iterate(v => flatData[i++] = v);
                sourceData = flatData;
            }

            // --- 2. Interact with WASM ---

            const exports = NDWasm.runtime.exports;
            let pixelDataPtr = 0;
            let formatPtr = 0;
            let resultPtr = 0;
            let encodedDataPtr = 0;

            const [height, width] = arrayToEncode.shape;
            const formatBytes = new TextEncoder().encode(format);

            try {
                // Allocate and copy pixel data to WASM.
                pixelDataPtr = exports.malloc(sourceData.byteLength);
                if (!pixelDataPtr) throw new Error("WASM malloc failed for pixel data.");
                new Uint8Array(exports.mem.buffer, pixelDataPtr, sourceData.byteLength).set(sourceData);

                // Allocate and copy format string to WASM.
                formatPtr = exports.malloc(formatBytes.length);
                if (!formatPtr) throw new Error("WASM malloc failed for format string.");
                new Uint8Array(exports.mem.buffer, formatPtr, formatBytes.length).set(formatBytes);

                // Call the WASM encoding function.
                resultPtr = exports.encode_image(pixelDataPtr, width, height, channels, quality, formatPtr, formatBytes.length);
                if (!resultPtr) {
                    console.error(`Failed to encode image to ${format} in WASM.`);
                    return null;
                }

                // Read the result structure [encoded_ptr, encoded_size].
                const resultView = new DataView(exports.mem.buffer, resultPtr, 8);
                encodedDataPtr = resultView.getUint32(0, true);
                const encodedDataSize = resultView.getUint32(4, true);

                if (!encodedDataPtr) {
                    console.error("WASM encode_image returned a null data pointer.");
                    return null;
                }

                // Create a copy of the encoded image data.
                const encodedData = new Uint8Array(exports.mem.buffer, encodedDataPtr, encodedDataSize).slice();

                return encodedData;

            } finally {
                // Free all allocated WASM memory.
                if (pixelDataPtr) exports.free(pixelDataPtr, sourceData.byteLength);
                if (formatPtr) exports.free(formatPtr, formatBytes.length);
                if (resultPtr) exports.free(resultPtr, 8);
                if (encodedDataPtr) {
                    const resultView = new DataView(exports.mem.buffer, resultPtr, 8);
                    const encodedDataSize = resultView.getUint32(4, true);
                    exports.free(encodedDataPtr, encodedDataSize);
                }
            }
        },

        /**
         * Encodes an NDArray into a PNG image.
         * This is a helper function that calls `encode` with `format: 'png'`.
         * @memberof NDWasmImage
         * @param {NDArray} ndarray - The input array. See `encode` for supported shapes and dtypes.
         * @returns {Uint8Array|null} A Uint8Array containing the binary data of the PNG image.
         * @example
         * const grayArr = ndarray.zeros([16, 16]);
         * const pngBytes = NDWasmImage.encodePng(grayArr);
         */
        encodePng(ndarray) {
            return this.encode(ndarray, { format: 'png' });
        },

        /**
         * Encodes an NDArray into a JPEG image.
         * This is a helper function that calls `encode` with `format: 'jpeg'`.
         * @memberof NDWasmImage
         * @param {NDArray} ndarray - The input array. See `encode` for supported shapes and dtypes.
         * @param {object} [options={}] - Encoding options.
         * @param {number} [options.quality=90] - The quality for JPEG encoding (1-100).
         * @returns {Uint8Array|null} A Uint8Array containing the binary data of the JPEG image.
         * @example
         * const floatArr = NDArray.random([20, 20, 3]);
         * const jpegBytes = NDWasmImage.encodeJpeg(floatArr, { quality: 85 });
         */
        encodeJpeg(ndarray, options = {}) {
            return this.encode(ndarray, { ...options, format: 'jpeg' });
        },

        /**
         * Converts a Uint8Array of binary data into a Base64 Data URL.
         * This is a utility function that runs purely in JavaScript.
         *
         * @memberof NDWasmImage
         * @param {Uint8Array} uint8array - The byte array to convert.
         * @param {string} [mimeType='image/png'] - The MIME type for the Data URL (e.g., 'image/jpeg').
         * @returns {string} The complete Data URL string.
         * @example
         * const pngBytes = NDWasmImage.encodePng(myNdarray);
         * const dataUrl = NDWasmImage.convertUint8ArrrayToDataurl(pngBytes, 'image/png');
         * // <img src={dataUrl} />
         */
        convertUint8ArrrayToDataurl(uint8array, mimeType = 'image/png') {
            // In a browser environment
            if (typeof window !== 'undefined' && typeof window.btoa === 'function') {
                let binary = '';
                const len = uint8array.byteLength;
                for (let i = 0; i < len; i++) {
                    binary += String.fromCharCode(uint8array[i]);
                }
                const base64 = window.btoa(binary);
                return `data:${mimeType};base64,${base64}`;
            }
            // In a Node.js environment
            if (typeof Buffer !== 'undefined') {
                const base64 = Buffer.from(uint8array).toString('base64');
                return `data:${mimeType};base64,${base64}`;
            }

            throw new Error("Unsupported environment: btoa or Buffer not available.");
        }
    };

    export default NDWasmImage;