/*
 * Copyright (c) 2024, Lucas Chollet <lucas.chollet@serenityos.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/ConstrainedStream.h>
#include <AK/MemoryStream.h>
#include <LibGfx/ImageFormats/JPEGXLCommon.h>
#include <LibGfx/ImageFormats/JPEGXLEntropyDecoder.h>

namespace Gfx {

namespace {
/// E.4.1 - Data stream
u8 icc_context(u64 i, u8 b1, u8 b2)
{
    u8 p1 = 0;
    u8 p2 = 0;

    if (i <= 128)
        return 0;

    if (b1 >= 'a' && b1 <= 'z')
        p1 = 0;
    else if (b1 >= 'A' && b1 <= 'Z')
        p1 = 0;
    else if (b1 >= '0' && b1 <= '9')
        p1 = 1;
    else if (b1 == '.' || b1 == ',')
        p1 = 1;
    else if (b1 <= 1)
        p1 = 2 + b1;
    else if (b1 > 1 && b1 < 16)
        p1 = 4;
    else if (b1 > 240 && b1 < 255)
        p1 = 5;
    else if (b1 == 255)
        p1 = 6;
    else
        p1 = 7;

    if (b2 >= 'a' && b2 <= 'z')
        p2 = 0;
    else if (b2 >= 'A' && b2 <= 'Z')
        p2 = 0;
    else if (b2 >= '0' && b2 <= '9')
        p2 = 1;
    else if (b2 == '.' || b2 == ',')
        p2 = 1;
    else if (b2 < 16)
        p2 = 2;
    else if (b2 > 240)
        p2 = 3;
    else
        p2 = 4;

    return 1 + p1 + p2 * 8;
}

ErrorOr<ByteBuffer> read_encoded_icc_stream(LittleEndianInputBitStream& stream)
{
    auto const enc_size = TRY(U64(stream));

    auto decoder = TRY(EntropyDecoder::create(stream, 41));

    ByteBuffer uncompressed_icc_stream {};
    TRY(uncompressed_icc_stream.try_resize(enc_size));
    for (u64 index = 0; index < enc_size; ++index) {
        auto const prev_byte = index > 0 ? uncompressed_icc_stream[index - 1] : 0u;
        auto const prev_prev_byte = index > 1 ? uncompressed_icc_stream[index - 2] : 0u;
        auto const context = icc_context(index, prev_byte, prev_prev_byte);
        uncompressed_icc_stream[index] = TRY(decoder.decode_hybrid_uint(stream, context));
    }

    TRY(decoder.ensure_end_state());
    return uncompressed_icc_stream;
}
///

/// E.4.2 - Encoded ICC stream
ErrorOr<u64> read_varint(Stream& stream)
{
    u64 value = 0;
    u8 shift = 0;
    while (1) {
        if (shift >= 56)
            return Error::from_string_literal("JPEGXLImageDecoderPlugin: Invalint shift value in varint");
        auto const b = TRY(stream.read_value<u8>());
        value += (b & 127) << shift;
        if (b <= 127)
            break;
        shift += 7;
    }
    return value;
}
///

/// E.4.3 - ICC header
ErrorOr<void> read_icc_header(Stream& data_stream, u32 output_size, ByteBuffer& out)
{
    u8 const header_size = min(128u, output_size);

    for (u8 i = 0; i < header_size; ++i) {
        auto const e = TRY(data_stream.read_value<u8>());
        u8 p = 0;

        if (i == 0 || i == 1 || i == 2 || i == 3) {
            // 'output_size[i]' means byte i of output_size encoded as an
            // unsigned 32-bit integer in big endian order
            BigEndian const output_size_as_be { static_cast<u32>(output_size) };
            p = bit_cast<u8 const*>(&output_size_as_be)[i];
        } else if (i == 8) {
            p = 4;
        } else if (i >= 12 && i <= 23) {
            auto s = "mntrRGB XYZ "sv;
            p = s[i - 12];
        } else if (i >= 36 && i <= 39) {
            auto s = "acsp"sv;
            p = s[i - 36];
        } else if ((i == 41 || i == 42) && out[40] == 'A') {
            p = 'P';
        } else if (i == 43 && out[40] == 'A') {
            p = 'L';
        } else if (i == 41 && out[40] == 'M') {
            p = 'S';
        } else if (i == 42 && out[40] == 'M') {
            p = 'F';
        } else if (i == 43 && out[40] == 'M') {
            p = 'T';
        } else if (i == 42 && out[40] == 'S' && out[41] == 'G') {
            p = 'I';
        } else if (i == 43 && out[40] == 'S' && out[41] == 'G') {
            p = 32;
        } else if (i == 42 && out[40] == 'S' && out[41] == 'U') {
            p = 'N';
        } else if (i == 43 && out[40] == 'S' && out[41] == 'U') {
            p = 'W';
        } else if (i == 70) {
            p = 246;
        } else if (i == 71) {
            p = 214;
        } else if (i == 73) {
            p = 1;
        } else if (i == 78) {
            p = 211;
        } else if (i == 79) {
            p = 45;
        } else if (i >= 80 && i < 84) {
            p = out[4 + i - 80];
        }

        out.append((p + e) & 255);
    }

    return {};
}
///

/// E.4.4 - ICC tag list

ErrorOr<void> append_as_u32_be(ByteBuffer& buffer, u32 value)
{
    BigEndian be_value { value };
    return buffer.try_append({ &be_value, sizeof(be_value) });
}
ErrorOr<void> read_tag_list(ConstrainedStream& command_stream, Stream& data_stream, ByteBuffer& out)
{
    auto const v = TRY(read_varint(command_stream));
    if (v == 0)
        return {};
    auto const num_tags = v - 1;

    TRY(append_as_u32_be(out, num_tags));
    u32 previous_tagstart = num_tags * 12 + 128;
    u32 previous_tagsize = 0;

    // Then, the decoder repeatedly reads a tag as specified by the following code until a tag with tagcode
    // equal to 0 is read or until the end of the command stream is reached.
    while (command_stream.remaining() > 0) {
        auto const command = TRY(command_stream.read_value<u8>());
        auto const tagcode = command & 63;
        if (tagcode == 0)
            return {};

        Array<u8, 4> tag;
        if (tagcode == 1) {
            TRY(data_stream.read_until_filled(tag));
        } else if (tagcode == 2) {
            tag = Array<u8, 4>::from_span("rTRC"sv.bytes());
        } else if (tagcode == 3) {
            tag = Array<u8, 4>::from_span("rXYZ"sv.bytes());
        } else if (tagcode >= 4 && tagcode < 21) {
            static constexpr auto strings = to_array(
                { "cprt"sv, "wtpt"sv, "bkpt"sv, "rXYZ"sv, "gXYZ"sv, "bXYZ"sv, "kXYZ"sv, "rTRC"sv, "gTRC"sv,
                    "bTRC"sv, "kTRC"sv, "chad"sv, "desc"sv, "chrm"sv, "dmnd"sv, "dmdd"sv, "lumi"sv });
            tag = Array<u8, 4>::from_span(strings[tagcode - 4].bytes());
        } else {
            return Error::from_string_literal("JPEGXLImageDecoderPlugin: Invalid tagcode in ICC profile");
        }

        auto tagstart = previous_tagstart + previous_tagsize;
        if ((command & 64) != 0)
            tagstart = TRY(read_varint(command_stream));

        u32 tagsize = previous_tagsize;
        if (tag == "rXYZ"sv.bytes() || tag == "gXYZ"sv.bytes() || tag == "bXYZ"sv.bytes()
            || tag == "kXYZ"sv.bytes() || tag == "wtpt"sv.bytes() || tag == "bkpt"sv.bytes()
            || tag == "lumi"sv.bytes())
            tagsize = 20;

        if ((command & 128) != 0)
            tagsize = TRY(read_varint(command_stream));
        previous_tagstart = tagstart;
        previous_tagsize = tagsize;

        // Write tag to output

        TRY(out.try_append(tag));
        TRY(append_as_u32_be(out, tagstart));
        TRY(append_as_u32_be(out, tagsize));
        if (tagcode == 2) {
            out.append("gTRC"sv.bytes());
            TRY(append_as_u32_be(out, tagstart));
            TRY(append_as_u32_be(out, tagsize));

            out.append("bTRC"sv.bytes());
            TRY(append_as_u32_be(out, tagstart));
            TRY(append_as_u32_be(out, tagsize));
        } else if (tagcode == 3) {
            out.append("gXYZ"sv.bytes());
            TRY(append_as_u32_be(out, tagstart + tagsize));
            TRY(append_as_u32_be(out, tagsize));

            out.append("bXYZ"sv.bytes());
            TRY(append_as_u32_be(out, tagstart + 2 * tagsize));
            TRY(append_as_u32_be(out, tagsize));
        }
    }

    return {};
}
///

/// E.4.5 - Main content
ErrorOr<void> shuffle(Bytes bytes, u8 width)
{
    auto temp = TRY(ByteBuffer::create_uninitialized(bytes.size()));

    u64 const height = (bytes.size() + width - 1) / width;

    u64 row_start = 0;
    u64 j = 0;
    for (size_t i = 0; i < bytes.size(); i++) {
        temp[i] = bytes[j];
        j += height;
        if (j >= bytes.size()) {
            ++row_start;
            j = row_start;
        }
    }

    temp.bytes().copy_to(bytes);
    return {};
}

ErrorOr<void> read_icc_main_content(ConstrainedStream& command_stream, Stream& data_stream, ByteBuffer& out)
{
    while (command_stream.remaining() > 0) {
        auto const command = TRY(command_stream.read_value<u8>());
        if (command == 1) {
            auto const num = TRY(read_varint(command_stream));
            auto bytes = TRY(out.get_bytes_for_writing(num));
            TRY(data_stream.read_until_filled(bytes));
        } else if (command == 2 or command == 3) {
            auto const num = TRY(read_varint(command_stream));
            auto bytes = TRY(out.get_bytes_for_writing(num));
            TRY(data_stream.read_until_filled(bytes));
            u8 width = (command == 2) ? 2 : 4;
            TRY(shuffle(bytes, width));
        } else if (command == 4) {
            u8 const flags = TRY(command_stream.read_value<u8>());
            u8 const width = (flags & 3) + 1;
            u8 const order = (flags & 12) >> 2;
            if (width == 3 || order == 3)
                return Error::from_string_literal("JPEGXLImageDecoderPlugin: Invalid width or order value");

            u64 stride = width;
            if ((flags & 16) != 0)
                stride = TRY(read_varint(command_stream));

            if (stride * 4 >= out.size() || stride < width)
                return Error::from_string_literal("JPEGXLImageDecoderPlugin: Invalid stride value");

            auto const num = TRY(read_varint(command_stream));
            ByteBuffer bytes;
            TRY(bytes.try_resize(num));
            TRY(data_stream.read_until_filled(bytes));
            if (width == 2 || width == 4)
                TRY(shuffle(bytes, width));

            for (u64 i = 0; i < num; i += width) {
                // NOTE: 0 <= order <= 2
                u8 const N = order + 1;
                Array<u32, 3> prev {};
                for (u8 j = 0; j < N; ++j) {
                    // "read u(width * 8) from the output ICC profile starting from
                    //  (stride * (j + 1)) bytes before the current output size,
                    //  interpreted as a big-endian unsigned integer of width bytes"
                    Array<u8, 4> bytes {};
                    for (u8 k = 0; k < width; ++k)
                        bytes[4 - width + k] = out[out.size() - stride * (j + 1) + k];
                    prev[j] = *bit_cast<BigEndian<u32> const*>(bytes.data());
                }

                u32 p;
                if (order == 0)
                    p = prev[0];
                else if (order == 1)
                    p = 2 * prev[0] - prev[1];
                else if (order == 2)
                    p = 3 * prev[0] - 3 * prev[1] + prev[2];

                for (u8 j = 0; j < width && i + j < num; ++j) {
                    u8 const val = (bytes[i + j] + (p >> (8 * (width - 1 - j)))) & 255;
                    TRY(out.try_append(val));
                }
            }
        } else if (command == 10) {
            TRY(out.try_append("XYZ "sv.bytes()));
            Array<u8, 4> constexpr zeros {};
            TRY(out.try_append(zeros));
            auto bytes = TRY(out.get_bytes_for_writing(12));
            TRY(data_stream.read_until_filled(bytes));
        } else if (command >= 16 and command < 24) {
            Array constexpr strings = { "XYZ "sv, "desc"sv, "text"sv, "mluc"sv, "para"sv, "curv"sv, "sf32"sv, "gbd "sv };
            TRY(out.try_append(strings[command - 16].bytes()));
            Array<u8, 4> constexpr zeros {};
            TRY(out.try_append(zeros));
        } else {
            return Error::from_string_literal("JPEGXLImageDecoderPlugin: Invalid command in ICC main context");
        }
    }

    return {};
}
///
}

/// E.4 - ICC profile
ErrorOr<ByteBuffer> read_icc(LittleEndianInputBitStream& stream)
{
    auto const encoded_icc = TRY(read_encoded_icc_stream(stream));

    FixedMemoryStream buffer(encoded_icc);
    auto const output_size = TRY(read_varint(buffer));
    auto const commands_size = TRY(read_varint(buffer));

    ConstrainedStream command_stream { MaybeOwned<Stream>(buffer), commands_size };

    auto const data_offset = buffer.offset() + commands_size;
    FixedMemoryStream data_stream(encoded_icc);
    TRY(data_stream.discard(data_offset));

    ByteBuffer out;
    TRY(out.try_ensure_capacity(output_size));

    TRY(read_icc_header(data_stream, output_size, out));

    if (output_size <= 128)
        return out;

    TRY(read_tag_list(command_stream, data_stream, out));

    TRY(read_icc_main_content(command_stream, data_stream, out));

    return out;
}
///

}