import { buildComponentData } from "./buildComponentData.js";
import { buildHuffmanTable } from "./buildHuffmanTable.js";
import { dctZigZag, decodeScan } from "./decodeScan.js";
import { JpegFrame } from "./JpegFrame.js";
import { JpegFrameComponent } from "./JpegFrameComponent.js";


/* -*- tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- /
/* vim: set shiftwidth=2 tabstop=2 autoindent cindent expandtab: */
/*
   Copyright 2026 Alex Goldring/ Company Named Ltd

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

// - The JPEG specification can be found in the ITU CCITT Recommendation T.81
//   (www.w3.org/Graphics/JPEG/itu-t81.pdf)
// - The JFIF specification can be found in the JPEG File Interchange Format
//   (www.w3.org/Graphics/JPEG/jfif3.pdf)
// - The Adobe Application-Specific JPEG markers in the Supporting the DCT Filters
//   in PostScript Level 2, Technical Note #5116
//   (partners.adobe.com/public/developer/en/ps/sdk/5116.DCT_Filter.pdf)

/**
 *
 * @param {number} a
 * @return {number}
 */
function clampTo8bit(a) {
    return a < 0 ? 0 : a > 255 ? 255 : a;
}

export class JpegImage {

    /**
     *
     * @param {Uint8Array} data
     */
    parse(data) {
        let offset = 0;

        function readUint16() {
            const value = (data[offset] << 8) | data[offset + 1];
            offset += 2;
            return value;
        }

        function readDataBlock() {
            const length = readUint16();
            const array = data.subarray(offset, offset + length - 2);
            offset += array.length;
            return array;
        }

        /**
         *
         * @param {JpegFrame} frame
         */
        function prepareComponents(frame) {
            // According to the JPEG standard, the sampling factor must be between 1 and 4
            // See https://github.com/libjpeg-turbo/libjpeg-turbo/blob/9abeff46d87bd201a952e276f3e4339556a403a3/libjpeg.txt#L1138-L1146
            let maxH = 1, maxV = 1;
            let component, componentId;
            for (componentId in frame.components) {
                if (frame.components.hasOwnProperty(componentId)) {
                    component = frame.components[componentId];
                    if (maxH < component.h) maxH = component.h;
                    if (maxV < component.v) maxV = component.v;
                }
            }
            const mcusPerLine = Math.ceil(frame.samplesPerLine / 8 / maxH);
            const mcusPerColumn = Math.ceil(frame.scanLines / 8 / maxV);
            for (componentId in frame.components) {
                if (frame.components.hasOwnProperty(componentId)) {
                    component = frame.components[componentId];
                    const blocksPerLine = Math.ceil(Math.ceil(frame.samplesPerLine / 8) * component.h / maxH);
                    const blocksPerColumn = Math.ceil(Math.ceil(frame.scanLines / 8) * component.v / maxV);
                    const blocksPerLineForMcu = mcusPerLine * component.h;
                    const blocksPerColumnForMcu = mcusPerColumn * component.v;
                    const blocksToAllocate = blocksPerColumnForMcu * blocksPerLineForMcu;
                    const blocks = [];

                    // Each block is a Int32Array of length 64 (4 x 64 = 256 bytes)

                    for (let i = 0; i < blocksPerColumnForMcu; i++) {
                        const row = [];
                        for (let j = 0; j < blocksPerLineForMcu; j++)
                            row.push(new Int32Array(64));
                        blocks.push(row);
                    }
                    component.blocksPerLine = blocksPerLine;
                    component.blocksPerColumn = blocksPerColumn;
                    component.blocks = blocks;
                }
            }
            frame.maxH = maxH;
            frame.maxV = maxV;
            frame.mcusPerLine = mcusPerLine;
            frame.mcusPerColumn = mcusPerColumn;
        }

        let jfif = null;
        let adobe = null;
        /**
         * @type {JpegFrame}
         */
        let frame;
        let resetInterval;
        /**
         *
         * @type {Int32Array[]}
         */
        const quantizationTables = []
        /**
         *
         * @type {JpegFrame[]}
         */
        const frames = [];
        const huffmanTablesAC = [], huffmanTablesDC = [];
        let fileMarker = readUint16();
        let malformedDataOffset = -1;

        /**
         *
         * @type {string[]}
         */
        this.comments = [];

        if (fileMarker !== 0xFFD8) { // SOI (Start of Image)
            throw new Error("SOI not found");
        }

        fileMarker = readUint16();
        while (fileMarker !== 0xFFD9) { // EOI (End of image)
            let i, j;
            switch (fileMarker) {
                case 0xFF00:
                    break;
                case 0xFFE0: // APP0 (Application Specific)
                case 0xFFE1: // APP1
                case 0xFFE2: // APP2
                case 0xFFE3: // APP3
                case 0xFFE4: // APP4
                case 0xFFE5: // APP5
                case 0xFFE6: // APP6
                case 0xFFE7: // APP7
                case 0xFFE8: // APP8
                case 0xFFE9: // APP9
                case 0xFFEA: // APP10
                case 0xFFEB: // APP11
                case 0xFFEC: // APP12
                case 0xFFEC: // APP12
                case 0xFFED: // APP13
                case 0xFFEE: // APP14
                case 0xFFEF: // APP15
                case 0xFFFE: // COM (Comment)
                    const appData = readDataBlock();

                    if (fileMarker === 0xFFFE) {
                        const comment = String.fromCharCode.apply(null, appData);
                        this.comments.push(comment);
                    }

                    if (fileMarker === 0xFFE0) {
                        if (appData[0] === 0x4A && appData[1] === 0x46 && appData[2] === 0x49 &&
                            appData[3] === 0x46 && appData[4] === 0) { // 'JFIF\x00'
                            jfif = {
                                version: { major: appData[5], minor: appData[6] },
                                densityUnits: appData[7],
                                xDensity: (appData[8] << 8) | appData[9],
                                yDensity: (appData[10] << 8) | appData[11],
                                thumbWidth: appData[12],
                                thumbHeight: appData[13],
                                thumbData: appData.subarray(14, 14 + 3 * appData[12] * appData[13])
                            };
                        }
                    }
                    // TODO APP1 - Exif
                    if (fileMarker === 0xFFE1) {
                        if (appData[0] === 0x45 &&
                            appData[1] === 0x78 &&
                            appData[2] === 0x69 &&
                            appData[3] === 0x66 &&
                            appData[4] === 0) { // 'EXIF\x00'
                            this.exifBuffer = appData.subarray(5, appData.length);
                        }
                    }

                    if (fileMarker === 0xFFEE) {
                        if (appData[0] === 0x41 && appData[1] === 0x64 && appData[2] === 0x6F &&
                            appData[3] === 0x62 && appData[4] === 0x65 && appData[5] === 0) { // 'Adobe\x00'
                            adobe = {
                                version: appData[6],
                                flags0: (appData[7] << 8) | appData[8],
                                flags1: (appData[9] << 8) | appData[10],
                                transformCode: appData[11]
                            };
                        }
                    }
                    break;

                case 0xFFDB: // DQT (Define Quantization Tables)
                    const quantizationTablesLength = readUint16();
                    const quantizationTablesEnd = quantizationTablesLength + offset - 2;
                    while (offset < quantizationTablesEnd) {

                        const quantizationTableSpec = data[offset++];
                        const tableData = new Int32Array(64);

                        if ((quantizationTableSpec >> 4) === 0) { // 8 bit values

                            for (j = 0; j < 64; j++) {
                                const z = dctZigZag[j];
                                tableData[z] = data[offset++];
                            }

                        } else if ((quantizationTableSpec >> 4) === 1) { //16 bit

                            for (j = 0; j < 64; j++) {
                                const z = dctZigZag[j];
                                tableData[z] = readUint16();
                            }

                        } else {
                            throw new Error("DQT: invalid table spec");
                        }

                        quantizationTables[quantizationTableSpec & 15] = tableData;
                    }
                    break;

                case 0xFFC0: // SOF0 (Start of Frame, Baseline DCT)
                case 0xFFC1: // SOF1 (Start of Frame, Extended DCT)
                case 0xFFC2: // SOF2 (Start of Frame, Progressive DCT)
                    readUint16(); // skip data length
                    frame = new JpegFrame();
                    frame.extended = (fileMarker === 0xFFC1);
                    frame.progressive = (fileMarker === 0xFFC2);
                    frame.precision = data[offset++];
                    frame.scanLines = readUint16();
                    frame.samplesPerLine = readUint16();
                    frame.components = {};
                    frame.componentsOrder = [];

                    let componentsCount = data[offset++], componentId;
                    for (i = 0; i < componentsCount; i++) {
                        componentId = data[offset];
                        const h = data[offset + 1] >> 4;
                        const v = data[offset + 1] & 15;
                        const qId = data[offset + 2];

                        if (h <= 0 || v <= 0) {
                            throw new Error('Invalid sampling factor, expected values above 0');
                        }

                        frame.componentsOrder.push(componentId);

                        const component = new JpegFrameComponent();
                        component.h = h;
                        component.v = v;
                        component.quantizationIdx = qId;

                        frame.components[componentId] = component;
                        offset += 3;
                    }
                    prepareComponents(frame);
                    frames.push(frame);
                    break;

                case 0xFFC4: // DHT (Define Huffman Tables)
                    const huffmanLength = readUint16();
                    const codeLengths = new Uint8Array(16);
                    for (i = 2; i < huffmanLength;) {
                        const huffmanTableSpec = data[offset++];
                        let codeLengthSum = 0;

                        for (j = 0; j < 16; j++, offset++) {
                            codeLengthSum += (codeLengths[j] = data[offset]);
                        }

                        const huffmanValues = new Uint8Array(codeLengthSum);

                        for (j = 0; j < codeLengthSum; j++, offset++) {
                            huffmanValues[j] = data[offset];
                        }

                        i += 17 + codeLengthSum;

                        const huffman_table = buildHuffmanTable(codeLengths, huffmanValues);

                        if ((huffmanTableSpec >> 4) === 0) {
                            huffmanTablesDC[huffmanTableSpec & 15] = huffman_table;
                        } else {
                            huffmanTablesAC[huffmanTableSpec & 15] = huffman_table;
                        }
                    }
                    break;

                case 0xFFDD: // DRI (Define Restart Interval)
                    readUint16(); // skip data length
                    resetInterval = readUint16();
                    break;

                case 0xFFDC: // Number of Lines marker
                    readUint16() // skip data length
                    readUint16() // Ignore this data since it represents the image height
                    break;

                case 0xFFDA: // SOS (Start of Scan)
                    const scanLength = readUint16();
                    const selectorsCount = data[offset++];
                    /**
                     *
                     * @type {JpegFrameComponent[]}
                     */
                    const components = [];
                    let component;
                    for (i = 0; i < selectorsCount; i++) {
                        component = frame.components[data[offset++]];
                        const tableSpec = data[offset++];
                        component.huffmanTableDC = huffmanTablesDC[tableSpec >> 4];
                        component.huffmanTableAC = huffmanTablesAC[tableSpec & 15];
                        components.push(component);
                    }
                    const spectralStart = data[offset++];
                    const spectralEnd = data[offset++];
                    const successiveApproximation = data[offset++];
                    const processed = decodeScan(
                        data,
                        offset,
                        frame,
                        components,
                        resetInterval,
                        spectralStart,
                        spectralEnd,
                        successiveApproximation >> 4,
                        successiveApproximation & 15,
                        this.opts.tolerantDecoding
                    );
                    offset += processed;
                    break;

                case 0xFFFF: // Fill bytes
                    if (data[offset] !== 0xFF) { // Avoid skipping a valid marker.
                        offset--;
                    }
                    break;
                default:
                    if (data[offset - 3] === 0xFF &&
                        data[offset - 2] >= 0xC0 && data[offset - 2] <= 0xFE) {
                        // could be incorrect encoding -- last 0xFF byte of the previous
                        // block was eaten by the encoder
                        offset -= 3;
                        break;
                    } else if (fileMarker === 0xE0 || fileMarker === 0xE1) {
                        // Recover from malformed APP1 markers popular in some phone models.
                        // See https://github.com/eugeneware/jpeg-js/issues/82
                        if (malformedDataOffset !== -1) {
                            throw new Error(`first unknown JPEG marker at offset ${malformedDataOffset.toString(16)}, second unknown JPEG marker ${fileMarker.toString(16)} at offset ${(offset - 1).toString(16)}`);
                        }
                        malformedDataOffset = offset - 1;
                        const nextOffset = readUint16();
                        if (data[offset + nextOffset - 2] === 0xFF) {
                            offset += nextOffset - 2;
                            break;
                        }
                    }
                    throw new Error("unknown JPEG marker " + fileMarker.toString(16));
            }
            fileMarker = readUint16();
        }
        if (frames.length !== 1)
            throw new Error("only single frame JPEGs supported");

        // set each frame's components quantization table
        for (let i = 0; i < frames.length; i++) {
            const cp = frames[i].components;
            for (let j in cp) {
                cp[j].quantizationTable = quantizationTables[cp[j].quantizationIdx];
                delete cp[j].quantizationIdx;
            }
        }

        this.width = frame.samplesPerLine;
        this.height = frame.scanLines;
        this.jfif = jfif;
        this.adobe = adobe;

        /**
         *
         * @type {{lines:Uint8Array[], scaleX:number, scaleY:number}[]}
         */
        this.components = [];

        for (let i = 0; i < frame.componentsOrder.length; i++) {
            const component = frame.components[frame.componentsOrder[i]];
            this.components.push({
                lines: buildComponentData(component),
                scaleX: component.h / frame.maxH,
                scaleY: component.v / frame.maxV
            });
        }
    }

    /**
     *
     * @param {number} width
     * @param {number} height
     * @return {Uint8Array}
     */
    getData(width, height) {
        const scaleX = this.width / width;
        const scaleY = this.height / height;

        let component1, component2, component3, component4;
        let component1Line, component2Line, component3Line, component4Line;
        let x, y;
        let offset = 0;
        let Y, Cb, Cr, K, C, M, Ye, R, G, B;
        let colorTransform;
        const dataLength = width * height * this.components.length;

        const data = new Uint8Array(dataLength);
        switch (this.components.length) {
            case 1:
                component1 = this.components[0];
                const lines = component1.lines;
                for (y = 0; y < height; y++) {
                    component1Line = lines[0 | (y * component1.scaleY * scaleY)];
                    for (x = 0; x < width; x++) {
                        Y = component1Line[0 | (x * component1.scaleX * scaleX)];

                        data[offset++] = Y;
                    }
                }
                break;
            case 2:
                // PDF might compress two component data in custom colorspace
                component1 = this.components[0];
                component2 = this.components[1];
                for (y = 0; y < height; y++) {
                    component1Line = component1.lines[0 | (y * component1.scaleY * scaleY)];
                    component2Line = component2.lines[0 | (y * component2.scaleY * scaleY)];
                    for (x = 0; x < width; x++) {
                        Y = component1Line[0 | (x * component1.scaleX * scaleX)];
                        data[offset++] = Y;
                        Y = component2Line[0 | (x * component2.scaleX * scaleX)];
                        data[offset++] = Y;
                    }
                }
                break;
            case 3:
                // The default transform for three components is true
                colorTransform = true;
                // The adobe transform marker overrides any previous setting
                if (this.adobe && this.adobe.transformCode)
                    colorTransform = true;
                else if (typeof this.opts.colorTransform !== 'undefined') {
                    colorTransform = !!this.opts.colorTransform;
                }

                component1 = this.components[0];
                component2 = this.components[1];
                component3 = this.components[2];
                for (y = 0; y < height; y++) {

                    component1Line = component1.lines[0 | (y * component1.scaleY * scaleY)];
                    component2Line = component2.lines[0 | (y * component2.scaleY * scaleY)];
                    component3Line = component3.lines[0 | (y * component3.scaleY * scaleY)];
                    for (x = 0; x < width; x++) {

                        if (!colorTransform) {
                            R = component1Line[0 | (x * component1.scaleX * scaleX)];
                            G = component2Line[0 | (x * component2.scaleX * scaleX)];
                            B = component3Line[0 | (x * component3.scaleX * scaleX)];
                        } else {
                            Y = component1Line[0 | (x * component1.scaleX * scaleX)];
                            Cb = component2Line[0 | (x * component2.scaleX * scaleX)];
                            Cr = component3Line[0 | (x * component3.scaleX * scaleX)];

                            R = clampTo8bit(Y + 1.402 * (Cr - 128));
                            G = clampTo8bit(Y - 0.3441363 * (Cb - 128) - 0.71413636 * (Cr - 128));
                            B = clampTo8bit(Y + 1.772 * (Cb - 128));
                        }

                        data[offset++] = R;
                        data[offset++] = G;
                        data[offset++] = B;
                    }
                }
                break;
            case 4:
                if (!this.adobe)
                    throw new Error('Unsupported color mode (4 components)');
                // The default transform for four components is false
                colorTransform = false;
                // The adobe transform marker overrides any previous setting
                if (this.adobe && this.adobe.transformCode)
                    colorTransform = true;
                else if (typeof this.opts.colorTransform !== 'undefined')
                    colorTransform = !!this.opts.colorTransform;

                component1 = this.components[0];
                component2 = this.components[1];
                component3 = this.components[2];
                component4 = this.components[3];
                for (y = 0; y < height; y++) {
                    const _y = y * scaleY;

                    component1Line = component1.lines[0 | (_y * component1.scaleY)];
                    component2Line = component2.lines[0 | (_y * component2.scaleY)];
                    component3Line = component3.lines[0 | (_y * component3.scaleY)];
                    component4Line = component4.lines[0 | (_y * component4.scaleY)];
                    for (x = 0; x < width; x++) {
                        const _x = x * scaleX;

                        if (!colorTransform) {
                            C = component1Line[0 | (_x * component1.scaleX)];
                            M = component2Line[0 | (_x * component2.scaleX)];
                            Ye = component3Line[0 | (_x * component3.scaleX)];
                            K = component4Line[0 | (_x * component4.scaleX)];
                        } else {
                            Y = component1Line[0 | (_x * component1.scaleX)];
                            Cb = component2Line[0 | (_x * component2.scaleX)];
                            Cr = component3Line[0 | (_x * component3.scaleX)];
                            K = component4Line[0 | (_x * component4.scaleX)];

                            C = 255 - clampTo8bit(Y + 1.402 * (Cr - 128));
                            M = 255 - clampTo8bit(Y - 0.3441363 * (Cb - 128) - 0.71413636 * (Cr - 128));
                            Ye = 255 - clampTo8bit(Y + 1.772 * (Cb - 128));
                        }

                        data[offset++] = 255 - C;
                        data[offset++] = 255 - M;
                        data[offset++] = 255 - Ye;
                        data[offset++] = 255 - K;
                    }
                }
                break;
            default:
                throw new Error('Unsupported color mode');
        }
        return data;
    }

    /**
     *
     * @param {{data:Uint8Array, width:number, height:number}} imageData
     * @param {boolean} [formatAsRGBA]
     */
    copyToImageData(imageData, formatAsRGBA = false) {
        const width = imageData.width, height = imageData.height;
        const imageDataArray = imageData.data;
        const data = this.getData(width, height);
        let i = 0, j = 0, x, y;
        let Y, K, C, M, R, G, B;
        switch (this.components.length) {
            case 1:
                for (y = 0; y < height; y++) {
                    for (x = 0; x < width; x++) {
                        Y = data[i++];

                        imageDataArray[j++] = Y;
                        imageDataArray[j++] = Y;
                        imageDataArray[j++] = Y;
                        if (formatAsRGBA) {
                            imageDataArray[j++] = 255;
                        }
                    }
                }
                break;
            case 3:
                for (y = 0; y < height; y++) {
                    for (x = 0; x < width; x++) {
                        R = data[i++];
                        G = data[i++];
                        B = data[i++];

                        imageDataArray[j++] = R;
                        imageDataArray[j++] = G;
                        imageDataArray[j++] = B;
                        if (formatAsRGBA) {
                            imageDataArray[j++] = 255;
                        }
                    }
                }
                break;
            case 4:
                for (y = 0; y < height; y++) {
                    for (x = 0; x < width; x++) {
                        C = data[i++];
                        M = data[i++];
                        Y = data[i++];
                        K = data[i++];

                        R = 255 - clampTo8bit(C * (1 - K / 255) + K);
                        G = 255 - clampTo8bit(M * (1 - K / 255) + K);
                        B = 255 - clampTo8bit(Y * (1 - K / 255) + K);

                        imageDataArray[j++] = R;
                        imageDataArray[j++] = G;
                        imageDataArray[j++] = B;
                        if (formatAsRGBA) {
                            imageDataArray[j++] = 255;
                        }
                    }
                }
                break;
            default:
                throw new Error('Unsupported color mode');
        }
    }
}

