/*

 * Copyright (c) 2024, Nico Weber <thakis@chromium.org>

 *

 * SPDX-License-Identifier: BSD-2-Clause

 */

#include <AK/BitStream.h>

#include <AK/Debug.h>

#include <AK/Enumerate.h>

#include <AK/MemoryStream.h>

#include <LibGfx/ICC/Profile.h>

#include <LibGfx/ImageFormats/ISOBMFF/JPEG2000Boxes.h>

#include <LibGfx/ImageFormats/ISOBMFF/Reader.h>

#include <LibGfx/ImageFormats/JPEG2000BitplaneDecoding.h>

#include <LibGfx/ImageFormats/JPEG2000InverseDiscreteWaveletTransform.h>

#include <LibGfx/ImageFormats/JPEG2000Loader.h>

#include <LibGfx/ImageFormats/JPEG2000ProgressionIterators.h>

#include <LibGfx/ImageFormats/JPEG2000TagTree.h>

#include <LibTextCodec/Decoder.h>

// Core coding system spec (.jp2 format): T-REC-T.800-201511-S!!PDF-E.pdf available here:

// https://www.itu.int/rec/dologin_pub.asp?lang=e&id=T-REC-T.800-201511-S!!PDF-E&type=items

// There is a useful example bitstream in the spec in:

// J.10 An example of decoding showing intermediate

// Extensions (.jpx format): T-REC-T.801-202106-S!!PDF-E.pdf available here:

// https://handle.itu.int/11.1002/1000/14666-en?locatt=format:pdf&auth

// rfc3745 lists the MIME type. It only mentions the jp2_id_string as magic number.

// A short overview of the JPEG2000 format:

//

// Image Decomposition

// -------------------

//

// 1. An image is first divided into independent tiles

// 2. Each tile is split into tile components (one each for R, G, B, A)

// 3. Each tile component undergoes Discrete Wavelet Transform (DWT)

//

// Resolution Levels and Subbands

// ------------------------------

//

// The DWT produces hierarchical resolution levels with these subbands:

// - Level 0: Single LL (Lowpass-Lowpass) subband

// - Level 1+: HL (Highpass-Lowpass), LH (Lowpass-Highpass), and HH (Highpass-Highpass) subbands

//

// Subband Layout:

// +-----+-----+----------+

// | LL0 | HL1 |          |

// +-----+-----+   HL2    |

// | LH1 | HH1 |          |

// +-----+-----+----------+

// |           |          |

// |    LH2    |    HH2   |

// |           |          |

// +-----------+----------+

//

// Precinct Structure

// ------------------

// - Precincts are rectangular regions that span all subbands within a resolution level

// - Typical size: 512k × 512k pixels

// - Most images contain only a single precinct due to this large size

// - "Precinct limited to a subband": portion of precinct covering one subband

//

// Layer System

// -----------

// - Coefficients' bitplanes can be stored separately

// - Groups of bitplanes form "layers"

// - For example, for an 8bpp image, layer 0 might contain the first two bitplanes, layer 1 the next two, etc.

// - Enables progressive refinement of image color resolution

//

// Codeblock Organization

// ----------------------

// - Each precinct is divided into codeblocks

// - A codeblock is the smallest coded unit in JPEG2000

// - Typical codeblock size: 64×64 pixels

// - Codeblocks store coefficient bitplanes from wavelet transformation

// - Independent arithmetic decoder contexts enable parallel decoding

// - A codeblock can be split into segments. A segment is a group of bytes

//   that are fed into the arithmetic decoder as one unit. Most files use one segment,

//   but the code block styles "termination on each coding pass" and

//   "selective arithmetic coding bypass" use multiple segments.

//

// Packets

// -------

// "All compressed image data representing a specific tile, layer, component, resolution level and precinct appears in the

//  codestream in a contiguous segment called a packet."

// A packet contains a packet header, and information about all codeblocks in the packet.

namespace Gfx {

// A JPEG2000 image can be stored in a codestream with markers, similar to a JPEG image,

// or in a JP2 file, which is a container format based on boxes similar to ISOBMFF.

// This is the marker for the codestream version.

// T.800 Annex A, Codestream syntax, A.2 Information in the marker segments and A.3 Construction of the codestream

static constexpr u8 marker_id_string[] = { 0xFF, 0x4F, 0xFF, 0x51 };

// This is the marker for the box version.

// T.800 Annex I, JP2 file format syntax, I.5.1 JPEG 2000 Signature box

static constexpr u8 jp2_id_string[] = { 0x00, 0x00, 0x00, 0x0C, 0x6A, 0x50, 0x20, 0x20, 0x0D, 0x0A, 0x87, 0x0A };

// Table A.2 – List of markers and marker segments

// "Delimiting markers and marker segments"

#define J2K_SOC 0xFF4F // "Start of codestream"

#define J2K_SOT 0xFF90 // "Start of tile-part"

#define J2K_SOD 0xFF93 // "Start of data"

#define J2K_EOC 0xFFD9 // "End of codestream"

// "Fixed information marker segments"

#define J2K_SIZ 0xFF51 // "Image and tile size"

// "Functional marker segments"

#define J2K_COD 0xFF52 // "Coding style default"

#define J2K_COC 0xFF53 // "Coding style component"

#define J2K_RGN 0xFF5E // "Region-of-interest"

#define J2K_QCD 0xFF5C // "Quantization default"

#define J2K_QCC 0xFF5D // "Quantization component"

#define J2K_POC 0xFF5F // "Progression order change"

// "Pointer marker segments"

#define J2K_TLM 0xFF55 // "Tile-part lengths"

#define J2K_PLM 0xFF57 // "Packet length, main header"

#define J2K_PLT 0xFF58 // "Packet length, tile-part header"

#define J2K_PPM 0xFF60 // "Packed packet headers, main header"

#define J2K_PPT 0xFF61 // "Packed packet headers, tile-part header"

// "In-bit-stream markers and marker segments"

#define J2K_SOP 0xFF91 // "Start of packet"

#define J2K_EPH 0xFF92 // "End of packet header"

// "Informational marker segments"

#define J2K_CRG 0xFF63 // "Component registration"

#define J2K_COM 0xFF64 // "Comment"

// A.4.2 Start of tile-part (SOT)

struct StartOfTilePart {

    // "Tile index. This number refers to the tiles in raster order starting at the number 0."

    u16 tile_index { 0 }; // "Isot" in spec.

    // "Length, in bytes, from the beginning of the first byte of this SOT marker segment of the tile-part to

    //  the end of the data of that tile-part. Figure A.16 shows this alignment. Only the last tile-part in the

    //  codestream may contain a 0 for Psot. If the Psot is 0, this tile-part is assumed to contain all data until

    //  the EOC marker."

    u32 tile_part_length { 0 }; // "Psot" in spec.

    // "Tile-part index. There is a specific order required for decoding tile-parts; this index denotes the order

    //  from 0. If there is only one tile-part for a tile, then this value is zero. The tile-parts of this tile shall

    //  appear in the codestream in this order, although not necessarily consecutively."

    u8 tile_part_index { 0 }; // "TPsot" in spec.

    // "Number of tile-parts of a tile in the codestream. Two values are allowed: the correct number of tile-

    //  parts for that tile and zero. A zero value indicates that the number of tile-parts of this tile is not

    //  specified in this tile-part.

    u8 number_of_tile_parts { 0 }; // "TNsot" in spec.

};

static ErrorOr<StartOfTilePart> read_start_of_tile_part(ReadonlyBytes data)

{

    FixedMemoryStream stream { data };

    StartOfTilePart sot;

    sot.tile_index = TRY(stream.read_value<BigEndian<u16>>());

    sot.tile_part_length = TRY(stream.read_value<BigEndian<u32>>());

    sot.tile_part_index = TRY(stream.read_value<u8>());

    sot.number_of_tile_parts = TRY(stream.read_value<u8>());

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: SOT marker segment: tile_index={}, tile_part_length={}, tile_part_index={}, number_of_tile_parts={}", sot.tile_index, sot.tile_part_length, sot.tile_part_index, sot.number_of_tile_parts);

    return sot;

}

// A.5.1 Image and tile size (SIZ)

struct ImageAndTileSize {

    // "Denotes capabilities that a decoder needs to properly decode the codestream."

    u16 needed_decoder_capabilities { 0 }; // "Rsiz" in spec.

    // "Width of the reference grid."

    u32 width { 0 }; // "Xsiz" in spec.

    // "Height of the reference grid."

    u32 height { 0 }; // "Ysiz" in spec.

    // "Horizontal offset from the origin of the reference grid to the left side of the image area."

    u32 x_offset { 0 }; // "XOsiz" in spec.

    // "Vertical offset from the origin of the reference grid to the top side of the image area."

    u32 y_offset { 0 }; // "YOsiz" in spec.

    // "Width of one reference tile with respect to the reference grid."

    u32 tile_width { 0 }; // "XTsiz" in spec.

    // "Height of one reference tile with respect to the reference grid."

    u32 tile_height { 0 }; // "YTsiz" in spec.

    // "Horizontal offset from the origin of the reference grid to the left side of the first tile."

    u32 tile_x_offset { 0 }; // "XTOsiz" in spec.

    // "Vertical offset from the origin of the reference grid to the top side of the first tile."

    u32 tile_y_offset { 0 }; // "YTOsiz" in spec.

    // "Csiz" isn't stored in this struct. It corresponds to `components.size()`.

    struct ComponentInformation {

        // "Precision (depth) in bits and sign of the ith component samples."

        u8 depth_and_sign { 0 }; // "Ssiz" in spec.

        // Table A.11 – Component Ssiz parameter

        u8 bit_depth() const { return (depth_and_sign & 0x7F) + 1; }

        bool is_signed() const { return depth_and_sign & 0x80; }

        // "Horizontal separation of a sample of the ith component with respect to the reference grid."

        u8 horizontal_separation { 0 }; // "XRsiz" in spec.

        // "Vertical separation of a sample of the ith component with respect to the reference grid."

        u8 vertical_separation { 0 }; // "YRsiz" in spec.

    };

    Vector<ComponentInformation> components;

    // (B-5)

    u32 number_of_x_tiles() const { return ceil_div(width - x_offset, tile_width); }

    u32 number_of_y_tiles() const { return ceil_div(height - y_offset, tile_height); }

    IntPoint tile_2d_index_from_1d_index(u32 tile_index) const

    {

        // (B-6)

        return { tile_index % number_of_x_tiles(), tile_index / number_of_x_tiles() };

    }

    IntRect reference_grid_coordinates_for_tile(IntPoint tile_2d_index) const

    {

        int p = tile_2d_index.x();

        int q = tile_2d_index.y();

        int tx0 = max(tile_x_offset + p * tile_width, x_offset);      // (B-7)

        int ty0 = max(tile_y_offset + q * tile_height, y_offset);     // (B-8)

        int tx1 = min(tile_x_offset + (p + 1) * tile_width, width);   // (B-9)

        int ty1 = min(tile_y_offset + (q + 1) * tile_height, height); // (B-10)

        return { tx0, ty0, tx1 - tx0, ty1 - ty0 };                    // (B-11)

    }

    IntRect reference_grid_coordinates_for_tile_component(IntRect tile_rect, int component_index) const

    {

        // (B-12)

        int tcx0 = ceil_div(tile_rect.left(), static_cast<int>(components[component_index].horizontal_separation));

        int tcx1 = ceil_div(tile_rect.right(), static_cast<int>(components[component_index].horizontal_separation));

        int tcy0 = ceil_div(tile_rect.top(), static_cast<int>(components[component_index].vertical_separation));

        int tcy1 = ceil_div(tile_rect.bottom(), static_cast<int>(components[component_index].vertical_separation));

        return { tcx0, tcy0, tcx1 - tcx0, tcy1 - tcy0 }; // (B-13)

    }

    IntRect reference_grid_coordinates_for_tile_component(IntPoint tile_2d_index, int component_index) const

    {

        auto tile_rect = reference_grid_coordinates_for_tile(tile_2d_index);

        return reference_grid_coordinates_for_tile_component(tile_rect, component_index);

    }

    IntRect reference_grid_coordinates_for_ll_band(IntRect tile_rect, int component_index, int r, int N_L) const

    {

        // B.5

        // (B-14)

        auto component_rect = reference_grid_coordinates_for_tile_component(tile_rect, component_index);

        int denominator = 1 << (N_L - r);

        int trx0 = ceil_div(component_rect.left(), denominator);

        int try0 = ceil_div(component_rect.top(), denominator);

        int trx1 = ceil_div(component_rect.right(), denominator);

        int try1 = ceil_div(component_rect.bottom(), denominator);

        return { trx0, try0, trx1 - trx0, try1 - try0 };

    }

    IntRect reference_grid_coordinates_for_sub_band(IntRect tile_rect, int component_index, int n_b, JPEG2000::SubBand sub_band) const

    {

        // B.5

        // Table B.1 – Quantities (xob, yob) for sub-band b

        int xob = 0;

        int yob = 0;

        if (sub_band == JPEG2000::SubBand::HorizontalHighpassVerticalLowpass || sub_band == JPEG2000::SubBand::HorizontalHighpassVerticalHighpass)

            xob = 1;

        if (sub_band == JPEG2000::SubBand::HorizontalLowpassVerticalHighpass || sub_band == JPEG2000::SubBand::HorizontalHighpassVerticalHighpass)

            yob = 1;

        VERIFY(n_b >= 1 || (n_b == 0 && sub_band == JPEG2000::SubBand::HorizontalLowpassVerticalLowpass));

        // If n_b is 0, `1 << (n_b - 1)` is undefined, but n_b is only 0 for the LL band, where xob and yob are 0 anyways.

        // So the value of o_scale doesn't matter in that case.

        int o_scale = 0;

        if (n_b > 0)

            o_scale = 1 << (n_b - 1);

        // (B-15)

        auto component_rect = reference_grid_coordinates_for_tile_component(tile_rect, component_index);

        int denominator = 1 << n_b;

        int tbx0 = ceil_div(component_rect.left() - o_scale * xob, denominator);

        int tby0 = ceil_div(component_rect.top() - o_scale * yob, denominator);

        int tbx1 = ceil_div(component_rect.right() - o_scale * xob, denominator);

        int tby1 = ceil_div(component_rect.bottom() - o_scale * yob, denominator);

        return { tbx0, tby0, tbx1 - tbx0, tby1 - tby0 };

    }

    IntRect reference_grid_coordinates_for_sub_band(IntPoint tile_2d_index, int component_index, int n_b, JPEG2000::SubBand sub_band) const

    {

        auto tile_rect = reference_grid_coordinates_for_tile(tile_2d_index);

        return reference_grid_coordinates_for_sub_band(tile_rect, component_index, n_b, sub_band);

    }

};

static ErrorOr<ImageAndTileSize> read_image_and_tile_size(ReadonlyBytes data)

{

    FixedMemoryStream stream { data };

    ImageAndTileSize siz;

    siz.needed_decoder_capabilities = TRY(stream.read_value<BigEndian<u16>>());

    siz.width = TRY(stream.read_value<BigEndian<u32>>());

    siz.height = TRY(stream.read_value<BigEndian<u32>>());

    siz.x_offset = TRY(stream.read_value<BigEndian<u32>>());

    siz.y_offset = TRY(stream.read_value<BigEndian<u32>>());

    siz.tile_width = TRY(stream.read_value<BigEndian<u32>>());

    siz.tile_height = TRY(stream.read_value<BigEndian<u32>>());

    siz.tile_x_offset = TRY(stream.read_value<BigEndian<u32>>());

    siz.tile_y_offset = TRY(stream.read_value<BigEndian<u32>>());

    u16 component_count = TRY(stream.read_value<BigEndian<u16>>()); // "Csiz" in spec.

    // Table A.9 – Image and tile size parameter values

    // Xsiz, Ysiz, XTsiz, YTsiz: 1 to 2^32-1.

    if (siz.width == 0 || siz.height == 0 || siz.tile_width == 0 || siz.tile_height == 0)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid image or tile size");

    // Ad-hoc: Limit image size to < 4 GiB.

    if (static_cast<u64>(siz.width) * siz.height > INT32_MAX)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Image is suspiciously large, not decoding");

    // CSiz: 1 to 16384.

    if (component_count < 1 || component_count > 16384)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid number of components");

    for (size_t i = 0; i < component_count; ++i) {

        ImageAndTileSize::ComponentInformation component;

        component.depth_and_sign = TRY(stream.read_value<u8>());

        if (component.bit_depth() > 38)

            return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid component depth");

        component.horizontal_separation = TRY(stream.read_value<u8>());

        component.vertical_separation = TRY(stream.read_value<u8>());

        siz.components.append(component);

    }

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: SIZ marker segment: needed_decoder_capabilities={}, width={}, height={}, x_offset={}, y_offset={}, tile_width={}, tile_height={}, tile_x_offset={}, tile_y_offset={}", siz.needed_decoder_capabilities, siz.width, siz.height, siz.x_offset, siz.y_offset, siz.tile_width, siz.tile_height, siz.tile_x_offset, siz.tile_y_offset);

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: SIZ marker segment: {} components:", component_count);

    for (auto [i, component] : enumerate(siz.components))

        dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: SIZ marker segment: component[{}]: is_signed={}, bit_depth={}, horizontal_separation={}, vertical_separation={}", i, component.is_signed(), component.bit_depth(), component.horizontal_separation, component.vertical_separation);

    return siz;

}

// Data shared by COD and COC marker segments

struct CodingStyleParameters {

    // Table A.15 – Coding style parameter values of the SPcod and SPcoc parameters

    // "Number of decomposition levels, NL, Zero implies no transformation."

    u8 number_of_decomposition_levels { 0 };

    u8 code_block_width_exponent { 0 };  // "xcb" in spec; 2 already added.

    u8 code_block_height_exponent { 0 }; // "ycb" in spec; 2 already added.

    u8 code_block_style { 0 };

    JPEG2000::Transformation transformation { JPEG2000::Transformation::Irreversible_9_7_Filter };

    // Table A.19 – Code-block style for the SPcod and SPcoc parameters

    bool uses_selective_arithmetic_coding_bypass() const { return code_block_style & 1; }

    bool reset_context_probabilities() const { return code_block_style & 2; }

    bool uses_termination_on_each_coding_pass() const { return code_block_style & 4; }

    bool uses_vertically_causal_context() const { return code_block_style & 8; }

    bool uses_predictable_termination() const { return code_block_style & 0x10; }

    bool uses_segmentation_symbols() const { return code_block_style & 0x20; }

    // If has_explicit_precinct_size is false, this contains the default { 15, 15 } number_of_decomposition_levels + 1 times.

    // If has_explicit_precinct_size is true, this contains number_of_decomposition_levels + 1 explicit values stored in the COD marker segment.

    struct PrecinctSize {

        u8 PPx { 0 };

        u8 PPy { 0 };

    };

    Vector<PrecinctSize> precinct_sizes;

};

static ErrorOr<CodingStyleParameters> read_coding_style_parameters(ReadonlyBytes data, StringView name, bool has_explicit_precinct_size)

{

    FixedMemoryStream stream { data };

    CodingStyleParameters parameters;

    parameters.number_of_decomposition_levels = TRY(stream.read_value<u8>());

    if (parameters.number_of_decomposition_levels > 32)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid number of decomposition levels");

    // Table A.18 – Width or height exponent of the code-blocks for the SPcod and SPcoc parameters

    u8 xcb = (TRY(stream.read_value<u8>()) & 0xF) + 2;

    u8 ycb = (TRY(stream.read_value<u8>()) & 0xF) + 2;

    if (xcb > 10 || ycb > 10 || xcb + ycb > 12)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid code block size");

    parameters.code_block_width_exponent = xcb;

    parameters.code_block_height_exponent = ycb;

    parameters.code_block_style = TRY(stream.read_value<u8>());

    if (parameters.code_block_style & 0xC0)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Reserved code block style bits set");

    // Table A.20 – Transformation for the SPcod and SPcoc parameters

    u8 transformation = TRY(stream.read_value<u8>());

    if (transformation > 1)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid transformation");

    parameters.transformation = transformation == 0 ? JPEG2000::Transformation::Irreversible_9_7_Filter : JPEG2000::Transformation::Reversible_5_3_Filter;

    if (has_explicit_precinct_size) {

        for (size_t i = 0; i < parameters.number_of_decomposition_levels + 1u; ++i) {

            u8 b = TRY(stream.read_value<u8>());

            // Table A.21 – Precinct width and height for the SPcod and SPcoc parameters

            CodingStyleParameters::PrecinctSize precinct_size;

            precinct_size.PPx = b & 0xF;

            precinct_size.PPy = b >> 4;

            if ((precinct_size.PPx == 0 || precinct_size.PPy == 0) && i > 0)

                return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid precinct size");

            parameters.precinct_sizes.append(precinct_size);

        }

    } else {

        for (size_t i = 0; i < parameters.number_of_decomposition_levels + 1u; ++i)

            parameters.precinct_sizes.append({ 15, 15 });

    }

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: number_of_decomposition_levels={}, code_block_width_exponent={}, code_block_height_exponent={}", name, parameters.number_of_decomposition_levels, parameters.code_block_width_exponent, parameters.code_block_height_exponent);

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: code_block_style={}, transformation={}", name, parameters.code_block_style, (int)parameters.transformation);

    if (has_explicit_precinct_size) {

        dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: {} explicit precinct sizes:", name, parameters.precinct_sizes.size());

        for (auto [i, precinct_size] : enumerate(parameters.precinct_sizes))

            dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: precinct_size[{}]: PPx={}, PPy={}", name, i, precinct_size.PPx, precinct_size.PPy);

    }

    return parameters;

}

// A.6.1 Coding style default (COD)

struct CodingStyleDefault {

    // Table A.13 – Coding style parameter values for the Scod parameter

    bool has_explicit_precinct_size { false };

    bool may_use_SOP_marker { false };

    bool shall_use_EPH_marker { false };

    // Table A.16 – Progression order for the SGcod, SPcoc, and Ppoc parameters

    // B.12 Progression order

    enum ProgressionOrder {

        LayerResolutionComponentPosition = 0,

        ResolutionLayerComponentPosition = 1,

        ResolutionPositionComponentLayer = 2,

        PositionComponentResolutionLayer = 3,

        ComponentPositionResolutionLayer = 4,

    };

    // Table A.17 – Multiple component transformation for the SGcod parameters

    enum MultipleComponentTransformationType {

        None = 0,

        MultipleComponentTransformationUsed = 1, // See Annex G

    };

    // Table A.14 – Coding style parameter values of the SGcod parameter

    ProgressionOrder progression_order { LayerResolutionComponentPosition };

    u16 number_of_layers { 0 };

    MultipleComponentTransformationType multiple_component_transformation_type { None };

    CodingStyleParameters parameters;

};

static ErrorOr<CodingStyleDefault> read_coding_style_default(ReadonlyBytes data)

{

    FixedMemoryStream stream { data };

    CodingStyleDefault cod;

    u8 Scod = TRY(stream.read_value<u8>());

    cod.has_explicit_precinct_size = Scod & 1;

    cod.may_use_SOP_marker = Scod & 2;

    cod.shall_use_EPH_marker = Scod & 4;

    u32 SGcod = TRY(stream.read_value<BigEndian<u32>>());

    u8 progression_order = SGcod >> 24;

    if (progression_order > 4)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid progression order");

    cod.progression_order = static_cast<CodingStyleDefault::ProgressionOrder>(progression_order);

    cod.number_of_layers = (SGcod >> 8) & 0xFFFF;

    if (cod.number_of_layers == 0)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid number of layers");

    u8 multiple_component_transformation_type = SGcod & 0xFF;

    if (multiple_component_transformation_type > 1)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid multiple component transformation type");

    cod.multiple_component_transformation_type = static_cast<CodingStyleDefault::MultipleComponentTransformationType>(multiple_component_transformation_type);

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: COD marker segment: has_explicit_precinct_size={}, may_use_SOP_marker={}, shall_use_EPH_marker={}", cod.has_explicit_precinct_size, cod.may_use_SOP_marker, cod.shall_use_EPH_marker);

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: COD marker segment: progression_order={}, number_of_layers={}, multiple_component_transformation_type={}", (int)cod.progression_order, cod.number_of_layers, (int)cod.multiple_component_transformation_type);

    cod.parameters = TRY(read_coding_style_parameters(data.slice(stream.offset()), "COD"sv, cod.has_explicit_precinct_size));

    return cod;

}

// A.6.2 Coding style component (COC)

struct CodingStyleComponent {

    u16 component_index { 0 }; // "Ccoc" in spec.

    // Table A.23 – Coding style parameter values for the Scoc parameter

    bool has_explicit_precinct_size { false }; // "Scoc" in spec.

    CodingStyleParameters parameters;

};

static ErrorOr<CodingStyleComponent> read_coding_style_component(ReadonlyBytes data, size_t number_of_components)

{

    FixedMemoryStream stream { data };

    // Table A.22 – Coding style component parameter values

    CodingStyleComponent coc;

    if (number_of_components < 257)

        coc.component_index = TRY(stream.read_value<u8>());

    else

        coc.component_index = TRY(stream.read_value<BigEndian<u16>>());

    u8 Scoc = TRY(stream.read_value<u8>());

    coc.has_explicit_precinct_size = Scoc & 1;

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: COC marker segment: component_index={}", coc.component_index);

    coc.parameters = TRY(read_coding_style_parameters(data.slice(TRY(stream.tell())), "COC"sv, coc.has_explicit_precinct_size));

    return coc;

}

// A.6.4 Quantization default (QCD)

struct QuantizationDefault {

    enum QuantizationStyle {

        NoQuantization = 0,

        ScalarDerived = 1,

        ScalarExpounded = 2,

    };

    QuantizationStyle quantization_style { NoQuantization };

    u8 number_of_guard_bits { 0 };

    struct ReversibleStepSize {

        u8 exponent { 0 };

    };

    struct IrreversibleStepSize {

        u16 mantissa { 0 };

        u8 exponent { 0 };

    };

    // Stores a Vector<ReversibleStepSize> if quantization_style is NoQuantization, and a Vector<IrreversibleStepSize> otherwise.

    // The size of the vector is >= 3*number_of_decomposition_levels + 1 if quantization_style is not ScalarDerived, and 1 otherwise.

    using StepSizeType = Variant<Empty, Vector<ReversibleStepSize>, Vector<IrreversibleStepSize>>;

    StepSizeType step_sizes;

};

static ErrorOr<QuantizationDefault> read_quantization_default(ReadonlyBytes data, StringView marker_name = "QCD"sv)

{

    FixedMemoryStream stream { data };

    QuantizationDefault qcd;

    u8 sqcd = TRY(stream.read_value<u8>());

    u8 quantization_style = sqcd & 0x1F;

    if (quantization_style > 2)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid quantization style");

    qcd.quantization_style = static_cast<QuantizationDefault::QuantizationStyle>(quantization_style);

    qcd.number_of_guard_bits = sqcd >> 5;

    qcd.step_sizes = TRY([&]() -> ErrorOr<QuantizationDefault::StepSizeType> {

        if (quantization_style == QuantizationDefault::NoQuantization) {

            // Table A.29 – Reversible step size values for the SPqcd and SPqcc parameters (reversible transform only)

            if (data.size() < 2)

                return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Not enough data for QCD marker segment");

            u8 number_of_decomposition_levels = (data.size() - 2) / 3;

            Vector<QuantizationDefault::ReversibleStepSize> reversible_step_sizes;

            for (size_t i = 0; i < 1u + 3u * number_of_decomposition_levels; ++i)

                reversible_step_sizes.append({ static_cast<u8>(TRY(stream.read_value<u8>()) >> 3) });

            return reversible_step_sizes;

        }

        // Table A.30 – Quantization values for the SPqcd and SPqcc parameters (irreversible transformation only)

        if (data.size() < 3)

            return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Not enough data for QCD marker segment");

        u8 number_of_decomposition_levels = 0;

        if (quantization_style == QuantizationDefault::ScalarExpounded)

            number_of_decomposition_levels = (data.size() - 3) / 6;

        Vector<QuantizationDefault::IrreversibleStepSize> irreversible_step_sizes;

        for (size_t i = 0; i < 1u + 3u * number_of_decomposition_levels; ++i) {

            u16 value = TRY(stream.read_value<BigEndian<u16>>());

            QuantizationDefault::IrreversibleStepSize step_size;

            step_size.mantissa = value & 0x7FF;

            step_size.exponent = value >> 11;

            irreversible_step_sizes.append(step_size);

        }

        return irreversible_step_sizes;

    }());

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: quantization_style={}, number_of_guard_bits={}", marker_name, (int)qcd.quantization_style, qcd.number_of_guard_bits);

    qcd.step_sizes.visit(

        [](Empty) { VERIFY_NOT_REACHED(); },

        [&](Vector<QuantizationDefault::ReversibleStepSize> const& step_sizes) {

            dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: {} step sizes:", marker_name, step_sizes.size());

            for (auto [i, step_size] : enumerate(step_sizes)) {

                dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: step_size[{}]: exponent={}", marker_name, i, step_size.exponent);

            }

        },

        [&](Vector<QuantizationDefault::IrreversibleStepSize> const& step_sizes) {

            dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: {} step sizes:", marker_name, step_sizes.size());

            for (auto [i, step_size] : enumerate(step_sizes)) {

                dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: {} marker segment: step_size[{}]: mantissa={}, exponent={}", marker_name, i, step_size.mantissa, step_size.exponent);

            }

        });

    return qcd;

}

// A.6.5 Quantization component (QCC)

struct QuantizationComponent {

    u16 component_index { 0 }; // "Cqcc" in spec.

    QuantizationDefault qcd;

};

static ErrorOr<QuantizationComponent> read_quantization_component(ReadonlyBytes data, size_t number_of_components)

{

    FixedMemoryStream stream { data };

    QuantizationComponent qcc;

    if (number_of_components < 257)

        qcc.component_index = TRY(stream.read_value<u8>());

    else

        qcc.component_index = TRY(stream.read_value<BigEndian<u16>>());

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: QCC marker segment: component_index={}", qcc.component_index);

    qcc.qcd = TRY(read_quantization_default(data.slice(TRY(stream.tell())), "QCC"sv));

    return qcc;

}

// A.9.2 Comment (COM)

struct Comment {

    enum CommentType {

        Binary = 0,

        ISO_IEC_8859_15 = 1,

    };

    CommentType type { Binary }; // "Rcom" in spec.

    ReadonlyBytes data;

};

static ErrorOr<Comment> read_comment(ReadonlyBytes data)

{

    FixedMemoryStream stream { data };

    Comment com;

    u16 comment_type = TRY(stream.read_value<BigEndian<u16>>());

    if (comment_type > 1)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Invalid comment type");

    com.type = static_cast<Comment::CommentType>(comment_type);

    com.data = data.slice(TRY(stream.tell()));

    dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: COM marker segment: comment_type={}, size()={}", (int)com.type, com.data.size());

    if (com.type == Comment::ISO_IEC_8859_15)

        dbgln_if(JPEG2000_DEBUG, "JPEG2000ImageDecoderPlugin: COM marker segment, ISO/IEC 8859-15 text: '{}'", TRY(TextCodec::decoder_for("ISO-8859-15"sv)->to_utf8(StringView { com.data })));

    return com;

}

struct TilePartData {

    StartOfTilePart sot;

    Vector<Comment> coms;

    ReadonlyBytes data;

};

struct DecodedCodeBlock {

    IntRect rect; // Confined to sub-band rect.

    // Transient state used to read packet headers.

    // B.10.4 Code-block inclusion

    bool is_included { false };

    // B.10.7.1 Single codeword segment

    // "Lblock is a code-block state variable. [...] The value of Lblock is initially set to three."

    u32 Lblock { 3 };

    // Becomes true when the first packet including this codeblock is read.

    bool has_been_included_in_previous_packet { false };

    // Data read from packet headers.

    // B.10.5 Zero bit-plane information

    // "the number of missing most significant bit-planes, P, may vary from code-block to code-block;

    //  these missing bit-planes are all taken to be zero."

    u32 p { 0 };

    struct Layer {

        struct Segment {

            ReadonlyBytes data;

            u32 index { 0 };

            int number_of_passes { 0 };

        };

        Vector<Segment, 1> segments;

    };

    Vector<Layer, 1> layers;

    u32 number_of_coding_passes() const

    {

        u32 total = 0;

        for (auto const& layer : layers)

            for (auto const& segment : layer.segments)

                total += segment.number_of_passes;

        return total;

    }

    u32 number_of_coding_passes_in_segment(u32 segment_index) const

    {

        u32 total = 0;

        for (auto const& layer : layers) {

            for (auto const& segment : layer.segments) {

                if (segment.index == segment_index)

                    total += segment.number_of_passes;

            }

        }

        return total;

    }

    Optional<u32> highest_segment_index() const

    {

        Optional<u32> highest_index;

        for (auto const& layer : layers) {

            for (auto const& segment : layer.segments)

                highest_index = max(highest_index.value_or(segment.index), segment.index);

        }

        return highest_index;

    }

    ErrorOr<Vector<ReadonlyBytes, 1>> segments_for_all_layers(ByteBuffer& maybe_storage) const

    {

        Vector<Vector<ReadonlyBytes, 1>, 1> all_segment_parts_for_segment;

        all_segment_parts_for_segment.resize(highest_segment_index().value_or(0) + 1);

        for (auto const& layer : layers)

            for (auto const& segment : layer.segments)

                TRY(all_segment_parts_for_segment[segment.index].try_append(segment.data));

        // Copy segments with multiple parts into consecutive storage.

        size_t total_scratch_size = 0;

        for (auto const& segment_parts : all_segment_parts_for_segment) {

            if (segment_parts.size() > 1) {

                for (auto const& segment_part : segment_parts)

                    total_scratch_size += segment_part.size();

            }

        }

        if (total_scratch_size > 0)

            maybe_storage = TRY(ByteBuffer::create_uninitialized(total_scratch_size));

        Vector<ReadonlyBytes, 1> all_segments;

        size_t scratch_offset = 0;

        for (auto& segment_parts : all_segment_parts_for_segment) {

            if (segment_parts.size() == 1) {

                TRY(all_segments.try_append(segment_parts[0]));

                continue;

            }

            auto start = scratch_offset;

            for (auto const& segment_part : segment_parts) {

                memcpy(maybe_storage.offset_pointer(scratch_offset), segment_part.data(), segment_part.size());

                scratch_offset += segment_part.size();

            }

            TRY(all_segments.try_append(maybe_storage.bytes().slice(start, scratch_offset - start)));

        }

        return all_segments;

    }

};

struct DecodedPrecinct {

    IntRect rect; // NOT confined to sub-band rect.

    int num_code_blocks_wide { 0 };

    int num_code_blocks_high { 0 };

    Vector<DecodedCodeBlock> code_blocks;

    // Transient state used to read packet headers.

    Optional<JPEG2000::TagTree> code_block_inclusion_tree;

    Optional<JPEG2000::TagTree> p_tree;

};

struct DecodedSubBand {

    IntRect rect;

    // These are the same for all three sub-bands at a given resolution level.

    int num_precincts_wide { 0 };

    int num_precincts_high { 0 };

    Vector<DecodedPrecinct> precincts;

    // Valid after bitplane decoding. rect.width() * rect.height() == coefficients.size().

    Vector<float> coefficients;

};

struct DecodedTileComponent {

    IntRect rect;

    DecodedSubBand nLL; // N_L LL in the spec, corresponds to resolution level 0.

    using DecodedSubBands = Array<DecodedSubBand, 3>; // Ordered HL, LH, HH.

    Vector<DecodedSubBands> decompositions;

    static constexpr Array SubBandOrder { JPEG2000::SubBand::HorizontalHighpassVerticalLowpass, JPEG2000::SubBand::HorizontalLowpassVerticalHighpass, JPEG2000::SubBand::HorizontalHighpassVerticalHighpass };

    // Valid after IDWT.

    Vector<float> samples;

};

struct TileData {

    // Data from codestream markers.

    Optional<CodingStyleDefault> cod;

    Vector<CodingStyleComponent> cocs;

    Optional<QuantizationDefault> qcd;

    Vector<QuantizationComponent> qccs;

    Vector<TilePartData> tile_parts;

    // Data used during decoding.

    IntRect rect;

    Vector<DecodedTileComponent> components;

    Vector<Vector<float>> channels;

    Vector<ImageAndTileSize::ComponentInformation> channel_information;

    // FIXME: This will have to move and be reorganized come POC support.

    OwnPtr<JPEG2000::ProgressionIterator> progression_iterator;

};

enum class ColorSpace {

    sRGB,

    Gray,

    CMYK,

    Unsupported,

};

struct JPEG2000LoadingContext {

    enum class State {

        NotDecoded = 0,

        DecodedImage,

        Error,

    };

    State state { State::NotDecoded };

    ReadonlyBytes codestream_data;

    size_t codestream_cursor { 0 };

    JPEG2000DecoderOptions options;

    Optional<ISOBMFF::JPEG2000ColorSpecificationBox const&> color_box; // This is always set for box-based files.

    Optional<ISOBMFF::JPEG2000PaletteBox const&> palette_box;

    Optional<ISOBMFF::JPEG2000ComponentMappingBox const&> component_mapping_box;

    Optional<ISOBMFF::JPEG2000ChannelDefinitionBox const&> channel_definition_box;

    IntSize size;

    ISOBMFF::BoxList boxes;

    // Data from marker segments:

    ImageAndTileSize siz;

    CodingStyleDefault cod;

    Vector<CodingStyleComponent> cocs;

    QuantizationDefault qcd;

    Vector<QuantizationComponent> qccs;

    Vector<Comment> coms;

    Vector<TileData> tiles;

    // Valid after headers have been decoded.

    // The awkward `color_space_error` is so that determine_color_space() can always succeed and

    // e.g. `file` can return data for JPEG2000s even if we can't decode the image data due to not

    // yet supporting its colorspace.

    ColorSpace color_space { ColorSpace::Unsupported };

    Optional<Error> color_space_error;

    // Valid once `state` is StateDecodedImage.

    RefPtr<Bitmap> bitmap;

    RefPtr<CMYKBitmap> cmyk_bitmap;

    CodingStyleParameters const& coding_style_parameters_for_component(TileData const& tile, size_t component_index) const

    {

        // Tile-part COC > Tile-part COD > Main COC > Main COD

        for (auto const& coc : tile.cocs) {

            if (coc.component_index == component_index)

                return coc.parameters;

        }

        if (tile.cod.has_value())

            return tile.cod->parameters;

        for (auto const& coc : cocs) {

            if (coc.component_index == component_index)

                return coc.parameters;

        }

        return cod.parameters;

    }

    QuantizationDefault const& quantization_parameters_for_component(TileData const& tile, size_t component_index)

    {

        // Tile-part QCC > Tile-part QCD > Main QCC > Main QCD

        for (auto const& qcc : tile.qccs) {

            if (qcc.component_index == component_index)

                return qcc.qcd;

        }

        if (tile.qcd.has_value())

            return tile.qcd.value();

        for (auto const& qcc : qccs) {

            if (qcc.component_index == component_index)

                return qcc.qcd;

        }

        return qcd;

    }

    ErrorOr<JPEG2000::ProgressionData> next_progression_data(TileData& tile) const

    {

        JPEG2000::ProgressionData progression_data;

        auto progression_data_has_packet = [&](JPEG2000::ProgressionData const& progression_data) {

            if (progression_data.resolution_level > coding_style_parameters_for_component(tile, progression_data.component).number_of_decomposition_levels)

                return false;

            // "It can happen that numprecincts is 0 for a particular tile-component and resolution level. When this happens, there are no

            //  packets for this tile-component and resolution level."

            // `num_precincts_wide` and `num_precincts_high` are the same for all sub-bands at a given resolution level, so it's

            // enough to only check the first.

            auto& component = tile.components[progression_data.component];

            auto& sub_band_data = progression_data.resolution_level == 0 ? component.nLL : component.decompositions[progression_data.resolution_level - 1][0];

            if (sub_band_data.num_precincts_wide == 0 || sub_band_data.num_precincts_high == 0)

                return false;

            return true;

        };

        do {

            if (!tile.progression_iterator->has_next())

                return Error::from_string_literal("JPEG2000ImageDecoderPlugin: No more progression orders but packets left");

            progression_data = tile.progression_iterator->next();

        } while (!progression_data_has_packet(progression_data));

        return progression_data;

    }

};

struct MarkerSegment {

    u16 marker;

    // OptionalNone for markers that don't have data.

    // For markers that do have data, this does not include the marker length data. (`data.size() + 2` is the value of the marker length field.)

    Optional<ReadonlyBytes> data;

};

static ErrorOr<u16> peek_marker(ReadonlyBytes data)

{

    if (2 > data.size())

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Not enough data for marker");

    return *reinterpret_cast<BigEndian<u16> const*>(data.data());

}

static ErrorOr<u16> peek_marker(JPEG2000LoadingContext& context)

{

    return peek_marker(context.codestream_data.slice(context.codestream_cursor));

}

static ErrorOr<MarkerSegment> read_marker_at_cursor(JPEG2000LoadingContext& context)

{

    u16 marker = TRY(peek_marker(context));

    // "All markers with the marker code between 0xFF30 and 0xFF3F have no marker segment parameters. They shall be skipped by the decoder."

    // "The SOC, SOD and EOC are delimiting markers not marker segments, and have no explicit length information or other parameters."

    bool is_marker_segment = !(marker >= 0xFF30 && marker <= 0xFF3F) && marker != J2K_SOC && marker != J2K_SOD && marker != J2K_EOC;

    MarkerSegment marker_segment;

    marker_segment.marker = marker;

    if (is_marker_segment) {

        if (context.codestream_cursor + 4 > context.codestream_data.size())

            return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Not enough data for marker segment length");

        u16 marker_length = *reinterpret_cast<BigEndian<u16> const*>(context.codestream_data.data() + context.codestream_cursor + 2);

        if (marker_length < 2)

            return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Marker segment length too small");

        if (context.codestream_cursor + 2 + marker_length > context.codestream_data.size())

            return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Not enough data for marker segment data");

        marker_segment.data = ReadonlyBytes { context.codestream_data.data() + context.codestream_cursor + 4, marker_length - 2u };

    }

    context.codestream_cursor += 2;

    if (is_marker_segment)

        context.codestream_cursor += 2 + marker_segment.data->size();

    return marker_segment;

}

static ErrorOr<void> parse_codestream_main_header(JPEG2000LoadingContext& context)

{

    // Figure A.3 – Construction of the main header

    // "Required as the first marker"

    auto marker = TRY(read_marker_at_cursor(context));

    if (marker.marker != J2K_SOC)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Expected SOC marker");

    // "Required as the second marker segment"

    marker = TRY(read_marker_at_cursor(context));

    if (marker.marker != J2K_SIZ)

        return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Expected SIZ marker");

    context.siz = TRY(read_image_and_tile_size(marker.data.value()));

    bool saw_COD_marker = false;

    bool saw_QCD_marker = false;

    while (true) {

        u16 marker = TRY(peek_marker(context));

        switch (marker) {

        case J2K_COD:

        case J2K_COC:

        case J2K_QCD:

        case J2K_QCC:

        case J2K_RGN:

        case J2K_POC:

        case J2K_PPM:

        case J2K_TLM:

        case J2K_PLM:

        case J2K_CRG:

        case J2K_COM: {

            auto marker = TRY(read_marker_at_cursor(context));

            if (marker.marker == J2K_COD) {

                if (saw_COD_marker)

                    return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Multiple COD markers in main header");

                context.cod = TRY(read_coding_style_default(marker.data.value()));

                saw_COD_marker = true;

            } else if (marker.marker == J2K_COC) {

                context.cocs.append(TRY(read_coding_style_component(marker.data.value(), context.siz.components.size())));

            } else if (marker.marker == J2K_QCD) {

                if (saw_QCD_marker)

                    return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Multiple QCD markers in main header");

                context.qcd = TRY(read_quantization_default(marker.data.value()));

                saw_QCD_marker = true;

            } else if (marker.marker == J2K_QCC) {

                context.qccs.append(TRY(read_quantization_component(marker.data.value(), context.siz.components.size())));

            } else if (marker.marker == J2K_COM) {

                context.coms.append(TRY(read_comment(marker.data.value())));

            } else if (marker.marker == J2K_TLM) {

                // TLM describes tile-part lengths, for random access. They can be ignored for now.

            } else if (marker.marker == J2K_PLM) {

                // PLM describes packet lengths, for random access. They can be ignored for now.

            } else {

                // FIXME: These are valid main header markers. Parse contents.

                dbgln("JPEG2000ImageDecoderPlugin: marker {:#04x} not yet implemented", marker.marker);

                return Error::from_string_literal("JPEG2000ImageDecoderPlugin: marker not yet implemented");

            }

            break;

        }

        case J2K_SOT: {

            // SOT terminates the main header.

            // A.4.2: "There shall be at least one SOT in a codestream."

            if (!saw_COD_marker)

                return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Required COD marker not present in main header");

            if (!saw_QCD_marker)

                return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Required QCD marker not present in main header");

            // A.6.4: "there is not necessarily a correspondence with the number of sub-bands present because the sub-bands

            //         can be truncated with no requirement to correct [the QCD] marker segment."

            size_t step_sizes_count = context.qcd.step_sizes.visit(

                [](Empty) -> size_t { VERIFY_NOT_REACHED(); },

                [](Vector<QuantizationDefault::ReversibleStepSize> const& step_sizes) { return step_sizes.size(); },

                [](Vector<QuantizationDefault::IrreversibleStepSize> const& step_sizes) { return step_sizes.size(); });

            // FIXME: What if number_of_decomposition_levels is in context.cocs and varies by component?

            if (context.qcd.quantization_style != QuantizationDefault::ScalarDerived && step_sizes_count < context.cod.parameters.number_of_decomposition_levels * 3u + 1u)

                return Error::from_string_literal("JPEG2000ImageDecoderPlugin: Not enough step sizes for number of decomposition levels");

            return {};