/*

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

 *

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

 */

#include <AK/BitStream.h>

#include <AK/Enumerate.h>

#include <LibGfx/ImageFormats/JBIG2Shared.h>

namespace Gfx::JBIG2 {

ErrorOr<TextRegionHuffmanTables> text_region_huffman_tables_from_flags(u16 huffman_flags, Vector<JBIG2::HuffmanTable const*> custom_tables)

{

    TextRegionHuffmanTables tables;

    u8 custom_table_index = 0;

    auto custom_table = [&custom_tables, &custom_table_index]() -> ErrorOr<JBIG2::HuffmanTable const*> {

        if (custom_table_index >= custom_tables.size())

            return Error::from_string_literal("JBIG2: Custom Huffman table index out of range");

        return custom_tables[custom_table_index++];

    };

    auto first_s_selection = (huffman_flags >> 0) & 0b11; // "SBHUFFFS" in spec.

    if (first_s_selection == 0)

        tables.first_s_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_6));

    else if (first_s_selection == 1)

        tables.first_s_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_7));

    else if (first_s_selection == 2)

        return Error::from_string_literal("JBIG2: Invalid first_s_table");

    else if (first_s_selection == 3)

        tables.first_s_table = TRY(custom_table());

    auto subsequent_s_selection = (huffman_flags >> 2) & 0b11; // "SBHUFFDS" in spec.

    if (subsequent_s_selection == 0)

        tables.subsequent_s_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_8));

    else if (subsequent_s_selection == 1)

        tables.subsequent_s_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_9));

    else if (subsequent_s_selection == 2)

        tables.subsequent_s_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_10));

    else if (subsequent_s_selection == 3)

        tables.subsequent_s_table = TRY(custom_table());

    auto delta_t_selection = (huffman_flags >> 4) & 0b11; // "SBHUFFDT" in spec.

    if (delta_t_selection == 0)

        tables.delta_t_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_11));

    else if (delta_t_selection == 1)

        tables.delta_t_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_12));

    else if (delta_t_selection == 2)

        tables.delta_t_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_13));

    else if (delta_t_selection == 3)

        tables.delta_t_table = TRY(custom_table());

    auto refinement_delta_width_selection = (huffman_flags >> 6) & 0b11; // "SBHUFFRDW" in spec.

    if (refinement_delta_width_selection == 0)

        tables.refinement_delta_width_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_14));

    else if (refinement_delta_width_selection == 1)

        tables.refinement_delta_width_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_15));

    else if (refinement_delta_width_selection == 2)

        return Error::from_string_literal("JBIG2: Invalid refinement_delta_width_table");

    else if (refinement_delta_width_selection == 3)

        tables.refinement_delta_width_table = TRY(custom_table());

    auto refinement_delta_height_selection = (huffman_flags >> 8) & 0b11; // "SBHUFFRDH" in spec.

    if (refinement_delta_height_selection == 0)

        tables.refinement_delta_height_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_14));

    else if (refinement_delta_height_selection == 1)

        tables.refinement_delta_height_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_15));

    else if (refinement_delta_height_selection == 2)

        return Error::from_string_literal("JBIG2: Invalid refinement_delta_height_table");

    else if (refinement_delta_height_selection == 3)

        tables.refinement_delta_height_table = TRY(custom_table());

    auto refinement_x_offset_selection = (huffman_flags >> 10) & 0b11; // "SBHUFFRDX" in spec.

    if (refinement_x_offset_selection == 0)

        tables.refinement_x_offset_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_14));

    else if (refinement_x_offset_selection == 1)

        tables.refinement_x_offset_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_15));

    else if (refinement_x_offset_selection == 2)

        return Error::from_string_literal("JBIG2: Invalid refinement_x_offset_table");

    else if (refinement_x_offset_selection == 3)

        tables.refinement_x_offset_table = TRY(custom_table());

    auto refinement_y_offset_selection = (huffman_flags >> 12) & 0b11; // "SBHUFFRDY" in spec.

    if (refinement_y_offset_selection == 0)

        tables.refinement_y_offset_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_14));

    else if (refinement_y_offset_selection == 1)

        tables.refinement_y_offset_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_15));

    else if (refinement_y_offset_selection == 2)

        return Error::from_string_literal("JBIG2: Invalid refinement_y_offset_table");

    else if (refinement_y_offset_selection == 3)

        tables.refinement_y_offset_table = TRY(custom_table());

    auto refinement_size_selection = (huffman_flags >> 14) & 0b1; // "SBHUFFRSIZE" in spec.

    if (refinement_size_selection == 0)

        tables.refinement_size_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_1));

    else if (refinement_size_selection == 1)

        tables.refinement_size_table = TRY(custom_table());

    if (custom_table_index != custom_tables.size())

        return Error::from_string_literal("JBIG2: Not all referred text region custom tables used");

    if (!tables.subsequent_s_table->has_oob_symbol())

        return Error::from_string_literal("JBIG2: Custom SBHUFFDS table must have OOB symbol");

    if (tables.first_s_table->has_oob_symbol()

        || tables.delta_t_table->has_oob_symbol()

        || tables.refinement_delta_width_table->has_oob_symbol()

        || tables.refinement_delta_height_table->has_oob_symbol()

        || tables.refinement_x_offset_table->has_oob_symbol()

        || tables.refinement_y_offset_table->has_oob_symbol()

        || tables.refinement_size_table->has_oob_symbol()) {

        return Error::from_string_literal("JBIG2: Custom text region Huffman tables must not have OOB symbol");

    }

    if (huffman_flags & 0x8000)

        return Error::from_string_literal("JBIG2: Invalid text region segment Huffman flags");

    return tables;

}

ErrorOr<SymbolDictionaryHuffmanTables> symbol_dictionary_huffman_tables_from_flags(u16 flags, Vector<JBIG2::HuffmanTable const*> custom_tables)

{

    // 7.4.2.1.1 Symbol dictionary flags

    SymbolDictionaryHuffmanTables tables;

    bool uses_huffman_encoding = (flags & 1) != 0; // "SDHUFF" in spec.

    u8 custom_table_index = 0;

    auto custom_table = [&custom_tables, &custom_table_index]() -> ErrorOr<JBIG2::HuffmanTable const*> {

        if (custom_table_index >= custom_tables.size())

            return Error::from_string_literal("JBIG2: Custom Huffman table index out of range");

        return custom_tables[custom_table_index++];

    };

    u8 huffman_table_selection_for_height_differences = (flags >> 2) & 0b11; // "SDHUFFDH" in spec.

    if (!uses_huffman_encoding && huffman_table_selection_for_height_differences != 0)

        return Error::from_string_literal("JBIG2: Invalid huffman_table_selection_for_height_differences");

    if (uses_huffman_encoding) {

        if (huffman_table_selection_for_height_differences == 0)

            tables.delta_height_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_4));

        else if (huffman_table_selection_for_height_differences == 1)

            tables.delta_height_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_5));

        else if (huffman_table_selection_for_height_differences == 2)

            return Error::from_string_literal("JBIG2: Invalid huffman_table_selection_for_height_differences");

        else if (huffman_table_selection_for_height_differences == 3)

            tables.delta_height_table = TRY(custom_table());

    }

    u8 huffman_table_selection_for_width_differences = (flags >> 4) & 0b11; // "SDHUFFDW" in spec.

    if (!uses_huffman_encoding && huffman_table_selection_for_width_differences != 0)

        return Error::from_string_literal("JBIG2: Invalid huffman_table_selection_for_width_differences");

    if (uses_huffman_encoding) {

        if (huffman_table_selection_for_width_differences == 0)

            tables.delta_width_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_2));

        else if (huffman_table_selection_for_width_differences == 1)

            tables.delta_width_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_3));

        else if (huffman_table_selection_for_width_differences == 2)

            return Error::from_string_literal("JBIG2: Invalid huffman_table_selection_for_width_differences");

        else if (huffman_table_selection_for_width_differences == 3)

            tables.delta_width_table = TRY(custom_table());

    }

    bool uses_user_supplied_size_table = (flags >> 6) & 1; // "SDHUFFBMSIZE" in spec.

    if (!uses_huffman_encoding && uses_user_supplied_size_table)

        return Error::from_string_literal("JBIG2: Invalid uses_user_supplied_size_table");

    if (uses_huffman_encoding) {

        if (!uses_user_supplied_size_table)

            tables.bitmap_size_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_1));

        else

            tables.bitmap_size_table = TRY(custom_table());

    }

    bool uses_user_supplied_aggregate_table = (flags >> 7) & 1; // "SDHUFFAGGINST" in spec.

    if (!uses_huffman_encoding && uses_user_supplied_aggregate_table)

        return Error::from_string_literal("JBIG2: Invalid uses_user_supplied_aggregate_table");

    if (uses_huffman_encoding) {

        if (!uses_user_supplied_aggregate_table)

            tables.number_of_symbol_instances_table = TRY(JBIG2::HuffmanTable::standard_huffman_table(JBIG2::HuffmanTable::StandardTable::B_1));

        else

            tables.number_of_symbol_instances_table = TRY(custom_table());

    }

    if (custom_table_index != custom_tables.size())

        return Error::from_string_literal("JBIG2: Not all referred symbol dictionary custom tables used");

    if (uses_huffman_encoding) {

        if (!tables.delta_width_table->has_oob_symbol())

            return Error::from_string_literal("JBIG2: Custom SDHUFFDW table must have OOB symbol");

        if (tables.delta_height_table->has_oob_symbol()

            || tables.bitmap_size_table->has_oob_symbol()

            || tables.number_of_symbol_instances_table->has_oob_symbol()) {

            return Error::from_string_literal("JBIG2: Custom symbol dictionary Huffman tables must not have OOB symbol");

        }

    }

    return tables;

}

ErrorOr<Vector<u32>> assign_huffman_codes(ReadonlyBytes code_lengths)

{

    // FIXME: Use shared huffman code, instead of using this algorithm from the spec.

    // B.3 Assigning the prefix codes

    // code_lengths is "PREFLEN" in spec, code_lengths.size is "NTEMP".

    Vector<u32> codes; // "CODES" in spec.

    TRY(codes.try_resize(code_lengths.size()));

    // "1) Build a histogram in the array LENCOUNT counting the number of times each prefix length value

    //     occurs in PREFLEN: LENCOUNT[I] is the number of times that the value I occurs in the array

    //     PREFLEN."

    Array<u32, 32> length_counts {}; // "LENCOUNT" in spec.

    for (auto length : code_lengths) {

        VERIFY(length < 32);

        length_counts[length]++;

    }

    // "2) Let LENMAX be the largest value for which LENCOUNT[LENMAX] > 0. Set:

    //         CURLEN = 1

    //         FIRSTCODE[0] = 0

    //         LENCOUNT[0] = 0"

    size_t highest_length_index = 0; // "LENMAX" in spec.

    for (auto const& [i, count] : enumerate(length_counts)) {

        if (count > 0)

            highest_length_index = i;

    }

    size_t current_length = 1;           // "CURLEN" in spec.

    Array<u32, 32> first_code_at_length; // "FIRSTCODE" in spec.

    first_code_at_length[0] = 0;

    length_counts[0] = 0;

    // "3) While CURLEN ≤ LENMAX, perform the following operations:"

    while (current_length <= highest_length_index) {

        // "a) Set:

        //         FIRSTCODE[CURLEN] = (FIRSTCODE[CURLEN – 1] + LENCOUNT[CURLEN – 1]) × 2

        //         CURCODE = FIRSTCODE[CURLEN]

        //         CURTEMP = 0"

        first_code_at_length[current_length] = (first_code_at_length[current_length - 1] + length_counts[current_length - 1]) * 2;

        u32 current_code = first_code_at_length[current_length]; // "CURCODE" in spec.

        size_t i = 0;                                            // "CURTEMP" in spec.

        // "b) While CURTEMP < NTEMP, perform the following operations:"

        while (i < code_lengths.size()) {

            // "i) If PREFLEN[CURTEMP] = CURLEN, then set:

            //         CODES[CURTEMP] = CURCODE

            //         CURCODE = CURCODE + 1"

            if (code_lengths[i] == current_length) {

                codes[i] = current_code;

                current_code++;

            }

            // "ii) Set CURTEMP = CURTEMP + 1"

            i++;

        }

        // "c) Set:

        //         CURLEN = CURLEN + 1"

        current_length++;

    }

    return codes;

}

ErrorOr<Vector<JBIG2::Code>> uniform_huffman_codes(u32 number_of_symbols, u32 code_length)

{

    Vector<u8> code_lengths;

    for (u32 i = 0; i < number_of_symbols; ++i)

        code_lengths.append(code_length);

    auto codes = TRY(JBIG2::assign_huffman_codes(code_lengths));

    Vector<JBIG2::Code> uniform_codes;

    for (auto const& [i, length] : enumerate(code_lengths)) {

        if (length == 0)

            continue;

        JBIG2::Code code { .prefix_length = length, .range_length = 0, .first_value = i, .code = codes[i] };

        uniform_codes.append(code);

    }

    return uniform_codes;

}

// Table B.1 – Standard Huffman table A

constexpr Array standard_huffman_table_A = {

    Code { 1, 4, 0, 0b0 },

    Code { 2, 8, 16, 0b10 },

    Code { 3, 16, 272, 0b110 },

    Code { 3, 32, 65808, 0b111 },

};

// Table B.2 – Standard Huffman table B

constexpr Array standard_huffman_table_B = {

    Code { 1, 0, 0, 0b0 },

    Code { 2, 0, 1, 0b10 },

    Code { 3, 0, 2, 0b110 },

    Code { 4, 3, 3, 0b1110 },

    Code { 5, 6, 11, 0b11110 },

    Code { 6, 32, 75, 0b111110 },

    Code { 6, 0, OptionalNone {}, 0b111111 },

};

// Table B.3 – Standard Huffman table C

constexpr Array standard_huffman_table_C = {

    Code { 8, 8, -256, 0b11111110 },

    Code { 1, 0, 0, 0b0 },

    Code { 2, 0, 1, 0b10 },

    Code { 3, 0, 2, 0b110 },

    Code { 4, 3, 3, 0b1110 },

    Code { 5, 6, 11, 0b11110 },

    Code { 8 | Code::LowerRangeBit, 32, -257, 0b11111111 },

    Code { 7, 32, 75, 0b1111110 },

    Code { 6, 0, OptionalNone {}, 0b111110 },

};

// Table B.4 – Standard Huffman table D

constexpr Array standard_huffman_table_D = {

    Code { 1, 0, 1, 0b0 },

    Code { 2, 0, 2, 0b10 },

    Code { 3, 0, 3, 0b110 },

    Code { 4, 3, 4, 0b1110 },

    Code { 5, 6, 12, 0b11110 },

    Code { 5, 32, 76, 0b11111 },

};

// Table B.5 – Standard Huffman table E

constexpr Array standard_huffman_table_E = {

    Code { 7, 8, -255, 0b1111110 },

    Code { 1, 0, 1, 0b0 },

    Code { 2, 0, 2, 0b10 },

    Code { 3, 0, 3, 0b110 },

    Code { 4, 3, 4, 0b1110 },

    Code { 5, 6, 12, 0b11110 },

    Code { 7 | Code::LowerRangeBit, 32, -256, 0b1111111 },

    Code { 6, 32, 76, 0b111110 },

};

// Table B.6 – Standard Huffman table F

constexpr Array standard_huffman_table_F = {

    Code { 5, 10, -2048, 0b11100 },

    Code { 4, 9, -1024, 0b1000 },

    Code { 4, 8, -512, 0b1001 },

    Code { 4, 7, -256, 0b1010 },

    Code { 5, 6, -128, 0b11101 },

    Code { 5, 5, -64, 0b11110 },

    Code { 4, 5, -32, 0b1011 },

    Code { 2, 7, 0, 0b00 },

    Code { 3, 7, 128, 0b010 },

    Code { 3, 8, 256, 0b011 },

    Code { 4, 9, 512, 0b1100 },

    Code { 4, 10, 1024, 0b1101 },

    Code { 6 | Code::LowerRangeBit, 32, -2049, 0b111110 },

    Code { 6, 32, 2048, 0b111111 },

};

// Table B.7 – Standard Huffman table G

constexpr Array standard_huffman_table_G = {

    Code { 4, 9, -1024, 0b1000 },

    Code { 3, 8, -512, 0b000 },

    Code { 4, 7, -256, 0b1001 },

    Code { 5, 6, -128, 0b11010 },

    Code { 5, 5, -64, 0b11011 },

    Code { 4, 5, -32, 0b1010 },

    Code { 4, 5, 0, 0b1011 },

    Code { 5, 5, 32, 0b11100 },

    Code { 5, 6, 64, 0b11101 },

    Code { 4, 7, 128, 0b1100 },

    Code { 3, 8, 256, 0b001 },

    Code { 3, 9, 512, 0b010 },

    Code { 3, 10, 1024, 0b011 },

    Code { 5 | Code::LowerRangeBit, 32, -1025, 0b11110 },

    Code { 5, 32, 2048, 0b11111 },

};

// Table B.8 – Standard Huffman table H

constexpr Array standard_huffman_table_H = {

    Code { 8, 3, -15, 0b11111100 },

    Code { 9, 1, -7, 0b111111100 },

    Code { 8, 1, -5, 0b11111101 },

    Code { 9, 0, -3, 0b111111101 },

    Code { 7, 0, -2, 0b1111100 },

    Code { 4, 0, -1, 0b1010 },

    Code { 2, 1, 0, 0b00 },

    Code { 5, 0, 2, 0b11010 },

    Code { 6, 0, 3, 0b111010 },

    Code { 3, 4, 4, 0b100 },

    Code { 6, 1, 20, 0b111011 },

    Code { 4, 4, 22, 0b1011 },

    Code { 4, 5, 38, 0b1100 },

    Code { 5, 6, 70, 0b11011 },

    Code { 5, 7, 134, 0b11100 },

    Code { 6, 7, 262, 0b111100 },

    Code { 7, 8, 390, 0b1111101 },

    Code { 6, 10, 646, 0b111101 },

    Code { 9 | Code::LowerRangeBit, 32, -16, 0b111111110 },

    Code { 9, 32, 1670, 0b111111111 },

    Code { 2, 0, OptionalNone {}, 0b01 },

};

// Table B.9 – Standard Huffman table I

constexpr Array standard_huffman_table_I = {

    Code { 8, 4, -31, 0b11111100 },

    Code { 9, 2, -15, 0b111111100 },

    Code { 8, 2, -11, 0b11111101 },

    Code { 9, 1, -7, 0b111111101 },

    Code { 7, 1, -5, 0b1111100 },

    Code { 4, 1, -3, 0b1010 },

    Code { 3, 1, -1, 0b010 },

    Code { 3, 1, 1, 0b011 },

    Code { 5, 1, 3, 0b11010 },

    Code { 6, 1, 5, 0b111010 },

    Code { 3, 5, 7, 0b100 },

    Code { 6, 2, 39, 0b111011 },

    Code { 4, 5, 43, 0b1011 },

    Code { 4, 6, 75, 0b1100 },

    Code { 5, 7, 139, 0b11011 },

    Code { 5, 8, 267, 0b11100 },

    Code { 6, 8, 523, 0b111100 },

    Code { 7, 9, 779, 0b1111101 },

    Code { 6, 11, 1291, 0b111101 },

    Code { 9 | Code::LowerRangeBit, 32, -32, 0b111111110 },

    Code { 9, 32, 3339, 0b111111111 },

    Code { 2, 0, OptionalNone {}, 0b00 },

};

// Table B.10 – Standard Huffman table J

constexpr Array standard_huffman_table_J = {

    Code { 7, 4, -21, 0b1111010 },

    Code { 8, 0, -5, 0b11111100 },

    Code { 7, 0, -4, 0b1111011 },

    Code { 5, 0, -3, 0b11000 },

    Code { 2, 2, -2, 0b00 },

    Code { 5, 0, 2, 0b11001 },

    Code { 6, 0, 3, 0b110110 },

    Code { 7, 0, 4, 0b1111100 },

    Code { 8, 0, 5, 0b11111101 },

    Code { 2, 6, 6, 0b01 },

    Code { 5, 5, 70, 0b11010 },

    Code { 6, 5, 102, 0b110111 },

    Code { 6, 6, 134, 0b111000 },

    Code { 6, 7, 198, 0b111001 },

    Code { 6, 8, 326, 0b111010 },

    Code { 6, 9, 582, 0b111011 },

    Code { 6, 10, 1094, 0b111100 },

    Code { 7, 11, 2118, 0b1111101 },

    Code { 8 | Code::LowerRangeBit, 32, -22, 0b11111110 },

    Code { 8, 32, 4166, 0b11111111 },

    Code { 2, 0, OptionalNone {}, 0b10 },

};

// Table B.11 – Standard Huffman table K

constexpr Array standard_huffman_table_K = {

    Code { 1, 0, 1, 0b0 },

    Code { 2, 1, 2, 0b10 },

    Code { 4, 0, 4, 0b1100 },

    Code { 4, 1, 5, 0b1101 },

    Code { 5, 1, 7, 0b11100 },

    Code { 5, 2, 9, 0b11101 },

    Code { 6, 2, 13, 0b111100 },

    Code { 7, 2, 17, 0b1111010 },

    Code { 7, 3, 21, 0b1111011 },

    Code { 7, 4, 29, 0b1111100 },

    Code { 7, 5, 45, 0b1111101 },

    Code { 7, 6, 77, 0b1111110 },

    Code { 7, 32, 141, 0b1111111 },

};

// Table B.12 – Standard Huffman table L

constexpr Array standard_huffman_table_L = {

    Code { 1, 0, 1, 0b0 },

    Code { 2, 0, 2, 0b10 },

    Code { 3, 1, 3, 0b110 },

    Code { 5, 0, 5, 0b11100 },

    Code { 5, 1, 6, 0b11101 },

    Code { 6, 1, 8, 0b111100 },

    Code { 7, 0, 10, 0b1111010 },

    Code { 7, 1, 11, 0b1111011 },

    Code { 7, 2, 13, 0b1111100 },

    Code { 7, 3, 17, 0b1111101 },

    Code { 7, 4, 25, 0b1111110 },

    Code { 8, 5, 41, 0b11111110 },

    Code { 8, 32, 73, 0b11111111 },

};

// Table B.13 – Standard Huffman table M

constexpr Array standard_huffman_table_M = {

    Code { 1, 0, 1, 0b0 },

    Code { 3, 0, 2, 0b100 },

    Code { 4, 0, 3, 0b1100 },

    Code { 5, 0, 4, 0b11100 },

    Code { 4, 1, 5, 0b1101 },

    Code { 3, 3, 7, 0b101 },

    Code { 6, 1, 15, 0b111010 },

    Code { 6, 2, 17, 0b111011 },

    Code { 6, 3, 21, 0b111100 },

    Code { 6, 4, 29, 0b111101 },

    Code { 6, 5, 45, 0b111110 },

    Code { 7, 6, 77, 0b1111110 },

    Code { 7, 32, 141, 0b1111111 },

};

// Table B.14 – Standard Huffman table N

constexpr Array standard_huffman_table_N = {

    Code { 3, 0, -2, 0b100 },

    Code { 3, 0, -1, 0b101 },

    Code { 1, 0, 0, 0b0 },

    Code { 3, 0, 1, 0b110 },

    Code { 3, 0, 2, 0b111 },

};

// Table B.15 – Standard Huffman table O

constexpr Array standard_huffman_table_O = {

    Code { 7, 4, -24, 0b1111100 },

    Code { 6, 2, -8, 0b111100 },

    Code { 5, 1, -4, 0b11100 },

    Code { 4, 0, -2, 0b1100 },

    Code { 3, 0, -1, 0b100 },

    Code { 1, 0, 0, 0b0 },

    Code { 3, 0, 1, 0b101 },

    Code { 4, 0, 2, 0b1101 },

    Code { 5, 1, 3, 0b11101 },

    Code { 6, 2, 5, 0b111101 },

    Code { 7, 4, 9, 0b1111101 },

    Code { 7 | Code::LowerRangeBit, 32, -25, 0b1111110 },

    Code { 7, 32, 25, 0b1111111 },

};

ErrorOr<HuffmanTable*> HuffmanTable::standard_huffman_table(StandardTable kind)

{

    switch (kind) {

    case StandardTable::B_1: {

        static HuffmanTable standard_table_A(standard_huffman_table_A);

        return &standard_table_A;

    }

    case StandardTable::B_2: {

        static HuffmanTable standard_table_B(standard_huffman_table_B, true);

        return &standard_table_B;

    }

    case StandardTable::B_3: {

        static HuffmanTable standard_table_C(standard_huffman_table_C, true);

        return &standard_table_C;

    }

    case StandardTable::B_4: {

        static HuffmanTable standard_table_D(standard_huffman_table_D);

        return &standard_table_D;

    }

    case StandardTable::B_5: {

        static HuffmanTable standard_table_E(standard_huffman_table_E);

        return &standard_table_E;

    }

    case StandardTable::B_6: {

        static HuffmanTable standard_table_F(standard_huffman_table_F);

        return &standard_table_F;

    }

    case StandardTable::B_7: {

        static HuffmanTable standard_table_G(standard_huffman_table_G);

        return &standard_table_G;

    }

    case StandardTable::B_8: {

        static HuffmanTable standard_table_H(standard_huffman_table_H, true);

        return &standard_table_H;

    }

    case StandardTable::B_9: {

        static HuffmanTable standard_table_I(standard_huffman_table_I, true);

        return &standard_table_I;

    }

    case StandardTable::B_10: {

        static HuffmanTable standard_table_J(standard_huffman_table_J, true);

        return &standard_table_J;

    }

    case StandardTable::B_11: {

        static HuffmanTable standard_table_K(standard_huffman_table_K);

        return &standard_table_K;

    }

    case StandardTable::B_12: {

        static HuffmanTable standard_table_L(standard_huffman_table_L);

        return &standard_table_L;

    }

    case StandardTable::B_13: {

        static HuffmanTable standard_table_M(standard_huffman_table_M);

        return &standard_table_M;

    }

    case StandardTable::B_14: {

        static HuffmanTable standard_table_N(standard_huffman_table_N);

        return &standard_table_N;

    }

    case StandardTable::B_15: {

        static HuffmanTable standard_table_O(standard_huffman_table_O);

        return &standard_table_O;

    }

    }

    VERIFY_NOT_REACHED();

}

ErrorOr<Optional<i32>> HuffmanTable::read_symbol_internal(BigEndianInputBitStream& stream) const

{

    // FIXME: Use an approach that doesn't require a full scan for every bit. See Compress::CanonicalCodes.

    u32 code_word = 0;

    u8 code_size = 0;

    while (true) {

        code_word = (code_word << 1) | TRY(stream.read_bit());

        code_size++;

        for (auto const& code : m_codes) {

            if ((code.prefix_length & ~Code::LowerRangeBit) == code_size && code.code == code_word) {

                if (!code.first_value.has_value())

                    return code.first_value; // OOB

                i32 value = 0; // "HTOFFSET" in spec.

                for (u8 i = 0; i < code.range_length; ++i)

                    value = (value << 1) | TRY(stream.read_bit());

                if (code.prefix_length & Code::LowerRangeBit)

                    return code.first_value.value() - value;

                return value + code.first_value.value();

            }

        }

    }

    VERIFY_NOT_REACHED();

}

ErrorOr<Optional<i32>> HuffmanTable::read_symbol(BigEndianInputBitStream& stream) const

{

    VERIFY(m_has_oob_symbol);

    return read_symbol_internal(stream);

}

ErrorOr<i32> HuffmanTable::read_symbol_non_oob(BigEndianInputBitStream& stream) const

{

    VERIFY(!m_has_oob_symbol);

    auto result = TRY(read_symbol_internal(stream));

    return result.value();

}

ErrorOr<void> HuffmanTable::write_symbol_internal(BigEndianOutputBitStream& stream, Optional<i32> value_or_oob) const

{

    // FIXME: Use an approach that doesn't require a full scan for every value,

    //        for example by handling OOB, lower range, and upper range first,

    //        and then binary searching the rest.

    for (auto const& code : m_codes) {

        if (value_or_oob.has_value() != code.first_value.has_value())

            continue;

        if (!value_or_oob.has_value()) {

            VERIFY(code.range_length == 0);

            return stream.write_bits(code.code, code.prefix_length);

        }

        auto first_value = code.first_value.value();

        auto value = value_or_oob.value();

        if (code.prefix_length & Code::LowerRangeBit) {

            VERIFY(code.range_length == 32);

            if (value > first_value)

                continue;

            TRY(stream.write_bits(code.code, code.prefix_length & ~Code::LowerRangeBit));

            return stream.write_bits(static_cast<u32>(first_value - value), code.range_length);

        }

        if (value < first_value || (code.range_length != 32 && value >= (first_value + (1 << code.range_length))))

            continue;

        TRY(stream.write_bits(code.code, code.prefix_length));

        return stream.write_bits(static_cast<u32>(value - first_value), code.range_length);

    }

    return Error::from_string_literal("JBIG2Writer: value not representable in this huffman table");

}

ErrorOr<void> HuffmanTable::write_symbol(BigEndianOutputBitStream& stream, Optional<i32> value) const

{

    VERIFY(m_has_oob_symbol);

    return write_symbol_internal(stream, value);

}

ErrorOr<void> HuffmanTable::write_symbol_non_oob(BigEndianOutputBitStream& stream, i32 value) const

{

    VERIFY(!m_has_oob_symbol);

    return write_symbol_internal(stream, value);

}

HuffmanTable::HuffmanTable(ReadonlySpan<Code> codes, bool has_oob_symbol)

    : m_codes(codes)

    , m_has_oob_symbol(has_oob_symbol)

{

}

}