/*
 * Copyright (c) 2023, Nicolas Ramz <nicolas.ramz@gmail.com>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/ByteReader.h>
#include <AK/Debug.h>
#include <AK/Endian.h>
#include <AK/FixedArray.h>
#include <AK/IntegralMath.h>
#include <LibCompress/PackBitsDecoder.h>
#include <LibGfx/FourCC.h>
#include <LibGfx/ImageFormats/ILBMLoader.h>
#include <LibRIFF/IFF.h>

namespace Gfx {

static constexpr size_t const ilbm_header_size = 12;

enum class CompressionType : u8 {
    None = 0,
    ByteRun = 1,
    __Count
};

enum class MaskType : u8 {
    None = 0,
    HasMask = 1,
    HasTransparentColor = 2,
    HasLasso = 3,
    __Count
};

enum class ViewportMode : u32 {
    EHB = 0x80,
    HAM = 0x800
};

enum class Format : u8 {
    // Amiga interleaved format
    ILBM = 0,
    // PC-DeluxePaint chunky format
    PBM = 1
};

AK_ENUM_BITWISE_OPERATORS(ViewportMode);

struct BMHDHeader {
    BigEndian<u16> width;
    BigEndian<u16> height;
    BigEndian<i16> x;
    BigEndian<i16> y;
    u8 planes;
    MaskType mask;
    CompressionType compression;
    u8 pad;
    BigEndian<u16> transparent_color;
    u8 x_aspect;
    u8 y_aspect;
    BigEndian<u16> page_width;
    BigEndian<u16> page_height;
};

static_assert(sizeof(BMHDHeader) == 20);

struct ILBMLoadingContext {
    enum class State {
        NotDecoded = 0,
        HeaderDecoded,
        BitmapDecoded
    };
    State state { State::NotDecoded };
    ReadonlyBytes data;

    // points to current chunk
    ReadonlyBytes chunks_cursor;

    // max number of bytes per plane row
    u16 pitch;

    ViewportMode viewport_mode;

    Vector<Color> color_table;

    // number of bits needed to describe current palette
    u8 cmap_bits;

    RefPtr<Gfx::Bitmap> bitmap;

    BMHDHeader bm_header;

    Format format;
};

static ErrorOr<void> decode_iff_ilbm_header(ILBMLoadingContext& context)
{
    if (context.state >= ILBMLoadingContext::State::HeaderDecoded)
        return {};

    if (context.data.size() < ilbm_header_size)
        return Error::from_string_literal("Missing IFF header");

    auto header_stream = FixedMemoryStream { context.data };
    auto header = TRY(IFF::FileHeader::read_from_stream(header_stream));
    if (header.magic() != "FORM"sv || (header.subformat != "ILBM"sv && header.subformat != "PBM "sv))
        return Error::from_string_literal("Invalid IFF-ILBM header");

    context.format = header.subformat == "ILBM" ? Format::ILBM : Format::PBM;

    return {};
}

static ErrorOr<Vector<Color>> decode_cmap_chunk(IFF::Chunk cmap_chunk)
{
    size_t const size = cmap_chunk.size() / 3;
    Vector<Color> color_table;
    TRY(color_table.try_ensure_capacity(size));

    for (size_t i = 0; i < size; ++i) {
        color_table.unchecked_append(Color(cmap_chunk[i * 3], cmap_chunk[(i * 3) + 1], cmap_chunk[(i * 3) + 2]));
    }

    return color_table;
}

static ErrorOr<RefPtr<Gfx::Bitmap>> chunky_to_bitmap(ILBMLoadingContext& context, ByteBuffer const& chunky)
{
    auto const width = context.bm_header.width;
    auto const height = context.bm_header.height;

    RefPtr<Gfx::Bitmap> bitmap = TRY(Bitmap::create(BitmapFormat::BGRA8888, { width, height }));

    dbgln_if(ILBM_DEBUG, "created Bitmap {}x{}", width, height);

    // - For 24bit pictures: the chunky buffer contains 3 bytes (R,G,B) per pixel
    // - For indexed colored pictures: chunky buffer contains a single byte per pixel
    u8 pixel_size = AK::max(1, context.bm_header.planes / 8);

    for (int row = 0; row < height; ++row) {
        // Keep color: in HAM mode, current color
        // may be based on previous color instead of coming from
        // the palette.
        Color color = Color::Black;
        for (int col = 0; col < width; col++) {
            size_t index = (width * row * pixel_size) + (col * pixel_size);
            if (context.bm_header.planes == 24) {
                color = Color(chunky[index], chunky[index + 1], chunky[index + 2]);
            } else if (chunky[index] < context.color_table.size()) {
                color = context.color_table[chunky[index]];
                if (context.bm_header.mask == MaskType::HasTransparentColor && chunky[index] == context.bm_header.transparent_color)
                    color = color.with_alpha(0);
            } else if (has_flag(context.viewport_mode, ViewportMode::HAM)) {
                // Get the control bit which will tell use how current pixel should be calculated
                u8 control = (chunky[index] >> context.cmap_bits) & 0x3;
                // Since we only have (cmap_bits - 2) bits to define the component,
                // we need to pad it to 8 bits.
                u8 component = (chunky[index] % context.color_table.size()) << (8 - context.cmap_bits);

                if (control == 1) {
                    color.set_blue(component);
                } else if (control == 2) {
                    color.set_red(component);
                } else {
                    color.set_green(component);
                }
            } else {
                return Error::from_string_literal("Color map index out of bounds but HAM bit not set");
            }
            bitmap->set_pixel(col, row, color);
        }
    }

    dbgln_if(ILBM_DEBUG, "filled Bitmap");

    return bitmap;
}

static ErrorOr<ByteBuffer> planar_to_chunky(ReadonlyBytes bitplanes, ILBMLoadingContext& context)
{
    dbgln_if(ILBM_DEBUG, "planar_to_chunky");
    u16 pitch = context.pitch;
    u16 width = context.bm_header.width;
    u16 height = context.bm_header.height;
    // mask is added as an extra plane
    u8 planes = context.bm_header.mask == MaskType::HasMask ? context.bm_header.planes + 1 : context.bm_header.planes;
    size_t buffer_size = static_cast<size_t>(width) * height;
    // If planes number is 24 we'll store R,G,B components so buffer needs to be 3 times width*height
    // otherwise we'll store a single 8bit index to the CMAP.
    if (planes == 24)
        buffer_size *= 3;

    auto chunky = TRY(ByteBuffer::create_zeroed(buffer_size));

    u8 const pixel_size = AK::max(1, planes / 8);

    for (u16 y = 0; y < height; y++) {
        size_t scanline = static_cast<size_t>(y) * width;
        for (u8 p = 0; p < planes; p++) {
            u8 const plane_mask = 1 << (p % 8);
            size_t offset_base = (pitch * planes * y) + (p * pitch);
            if (offset_base + pitch > bitplanes.size())
                return Error::from_string_literal("Malformed bitplane data");

            for (u16 i = 0; i < pitch; i++) {
                u8 bit = bitplanes[offset_base + i];
                u8 rgb_shift = p / 8;

                // Some encoders don't pad bytes rows with 0: make sure we stop
                // when enough data for current bitplane row has been read
                for (u8 b = 0; b < 8 && (i * 8) + b < width; b++) {
                    u8 mask = 1 << (7 - b);
                    // get current plane: simply skip mask plane for now
                    if (bit & mask && p < context.bm_header.planes) {
                        u16 x = (i * 8) + b;
                        size_t offset = (scanline * pixel_size) + (x * pixel_size) + rgb_shift;
                        // Only throw an error if we would actually attempt to write
                        // outside of the chunky buffer. Some apps like PPaint produce
                        // malformed bitplane data but files are still accepted by most readers
                        // since they do not cause writing past the chunky buffer.
                        if (offset >= chunky.size())
                            return Error::from_string_literal("Malformed bitplane data");

                        chunky[offset] |= plane_mask;
                    }
                }
            }
        }
    }

    dbgln_if(ILBM_DEBUG, "planar_to_chunky: end");

    return chunky;
}

static ErrorOr<ByteBuffer> uncompress_byte_run(ReadonlyBytes data, ILBMLoadingContext& context)
{
    auto length = data.size();
    dbgln_if(ILBM_DEBUG, "uncompress_byte_run pitch={} size={}", context.pitch, data.size());

    size_t plane_data_size = context.pitch * context.bm_header.height * context.bm_header.planes;

    // The mask is encoded as an extra bitplane but is not counted in the bm_header planes
    if (context.bm_header.mask == MaskType::HasMask)
        plane_data_size += context.pitch * context.bm_header.height;

    // The maximum run length of this compression method is 127 bytes, so the uncompressed size
    // cannot be more than 127 times the size of the chunk we are decompressing.
    if (plane_data_size > NumericLimits<u32>::max() || ceil_div(plane_data_size, 127ul) > length)
        return Error::from_string_literal("Uncompressed data size too large");

    auto plane_data = TRY(Compress::PackBits::decode_all(data, plane_data_size));

    return plane_data;
}

static ErrorOr<void> extend_ehb_palette(ILBMLoadingContext& context)
{
    dbgln_if(ILBM_DEBUG, "need to extend palette");
    for (size_t i = 0; i < 32; ++i) {
        auto const color = context.color_table[i];
        TRY(context.color_table.try_append(color.darkened()));
    }

    return {};
}

static ErrorOr<void> reduce_ham_palette(ILBMLoadingContext& context)
{
    u8 bits = context.cmap_bits;

    dbgln_if(ILBM_DEBUG, "reduce palette planes={} bits={}", context.bm_header.planes, context.cmap_bits);

    if (bits > context.bm_header.planes) {
        dbgln_if(ILBM_DEBUG, "need to reduce palette");
        bits -= (bits - context.bm_header.planes) + 2;
        // bits shouldn't theorically be less than 4 bits in HAM mode.
        if (bits < 4)
            return Error::from_string_literal("Error while reducing CMAP for HAM: bits too small");

        context.color_table.resize((context.color_table.size() >> bits));
        context.cmap_bits = bits;
    }

    return {};
}

static ErrorOr<void> decode_body_chunk(IFF::Chunk body_chunk, ILBMLoadingContext& context)
{
    dbgln_if(ILBM_DEBUG, "decode_body_chunk {}", body_chunk.size());

    ByteBuffer pixel_data;

    if (context.bm_header.compression == CompressionType::ByteRun) {
        auto plane_data = TRY(uncompress_byte_run(body_chunk.data(), context));
        if (context.format == Format::ILBM)
            pixel_data = TRY(planar_to_chunky(plane_data, context));
        else
            pixel_data = plane_data;
    } else {
        if (context.format == Format::ILBM)
            pixel_data = TRY(planar_to_chunky(body_chunk.data(), context));
        else
            pixel_data = TRY(ByteBuffer::copy(body_chunk.data().data(), body_chunk.size()));
    }

    // Some files already have 64 colors defined in the palette,
    // maybe for upward compatibility with 256 colors software/hardware.
    // DPaint 4 & previous files only have 32 colors so the
    // palette needs to be extended only for these files.
    if (has_flag(context.viewport_mode, ViewportMode::EHB) && context.color_table.size() < 64) {
        TRY(extend_ehb_palette(context));
    } else if (has_flag(context.viewport_mode, ViewportMode::HAM)) {
        TRY(reduce_ham_palette(context));
    }

    context.bitmap = TRY(chunky_to_bitmap(context, pixel_data));

    return {};
}

static ErrorOr<void> decode_iff_chunks(ILBMLoadingContext& context)
{
    auto& chunks = context.chunks_cursor;

    dbgln_if(ILBM_DEBUG, "decode_iff_chunks");

    while (!chunks.is_empty()) {
        auto chunk = TRY(IFF::Chunk::decode_and_advance(chunks));
        if (chunk.id() == "CMAP"sv) {
            // Some files (HAM mainly) have CMAP chunks larger than the planes they advertise: I'm not sure
            // why but we should not return an error in this case.

            context.color_table = TRY(decode_cmap_chunk(chunk));
            context.cmap_bits = AK::ceil_log2(context.color_table.size());
        } else if (chunk.id() == "BODY"sv) {
            if (context.color_table.is_empty() && context.bm_header.planes != 24)
                return Error::from_string_literal("Decoding indexed BODY chunk without a color map is not currently supported");

            // Apparently 32bit ilbm files exist: but I wasn't able to find any,
            // nor is it documented anywhere, so let's make it clear it's not supported.
            if (context.bm_header.planes != 24 && context.bm_header.planes > 8)
                return Error::from_string_literal("Invalid number of bitplanes");

            TRY(decode_body_chunk(chunk, context));
            context.state = ILBMLoadingContext::State::BitmapDecoded;
        } else if (chunk.id() == "CRNG"sv) {
            dbgln_if(ILBM_DEBUG, "Chunk:CRNG");
        } else if (chunk.id() == "CAMG"sv) {
            context.viewport_mode = static_cast<ViewportMode>(AK::convert_between_host_and_big_endian(ByteReader::load32(chunk.data().data())));
            dbgln_if(ILBM_DEBUG, "Chunk:CAMG, Viewport={}, EHB={}, HAM={}", (u32)context.viewport_mode, has_flag(context.viewport_mode, ViewportMode::EHB), has_flag(context.viewport_mode, ViewportMode::HAM));
        }
    }

    if (context.state != ILBMLoadingContext::State::BitmapDecoded)
        return Error::from_string_literal("Missing body chunk");

    return {};
}

static ErrorOr<void> decode_bmhd_chunk(ILBMLoadingContext& context)
{
    context.chunks_cursor = context.data.slice(sizeof(IFF::FileHeader));
    auto first_chunk = TRY(IFF::Chunk::decode_and_advance(context.chunks_cursor));

    if (first_chunk.id() != "BMHD"sv)
        return Error::from_string_literal("IFFImageDecoderPlugin: Invalid chunk type, expected BMHD");

    if (first_chunk.size() < sizeof(BMHDHeader))
        return Error::from_string_literal("IFFImageDecoderPlugin: Not enough data for header chunk");

    context.bm_header = *bit_cast<BMHDHeader const*>(first_chunk.data().data());

    if (context.bm_header.mask >= MaskType::__Count)
        return Error::from_string_literal("IFFImageDecoderPlugin: Unsupported mask type");

    if (context.bm_header.compression >= CompressionType::__Count)
        return Error::from_string_literal("IFFImageDecoderPlugin: Unsupported compression type");

    context.pitch = ceil_div((u16)context.bm_header.width, (u16)16) * 2;

    context.state = ILBMLoadingContext::State::HeaderDecoded;

    dbgln_if(ILBM_DEBUG, "IFFImageDecoderPlugin: BMHD: {}x{} ({},{}), p={}, m={}, c={}",
        context.bm_header.width,
        context.bm_header.height,
        context.bm_header.x,
        context.bm_header.y,
        context.bm_header.planes,
        to_underlying(context.bm_header.mask),
        to_underlying(context.bm_header.compression));

    return {};
}

ILBMImageDecoderPlugin::ILBMImageDecoderPlugin(ReadonlyBytes data, NonnullOwnPtr<ILBMLoadingContext> context)
    : m_context(move(context))
{
    m_context->data = data;
}

ILBMImageDecoderPlugin::~ILBMImageDecoderPlugin() = default;

IntSize ILBMImageDecoderPlugin::size()
{
    return IntSize { m_context->bm_header.width, m_context->bm_header.height };
}

bool ILBMImageDecoderPlugin::sniff(ReadonlyBytes data)
{
    ILBMLoadingContext context;
    context.data = data;

    return !decode_iff_ilbm_header(context).is_error();
}

ErrorOr<NonnullOwnPtr<ImageDecoderPlugin>> ILBMImageDecoderPlugin::create(ReadonlyBytes data)
{
    auto context = TRY(try_make<ILBMLoadingContext>());
    auto plugin = TRY(adopt_nonnull_own_or_enomem(new (nothrow) ILBMImageDecoderPlugin(data, move(context))));
    TRY(decode_iff_ilbm_header(*plugin->m_context));
    TRY(decode_bmhd_chunk(*plugin->m_context));
    return plugin;
}

ErrorOr<ImageFrameDescriptor> ILBMImageDecoderPlugin::frame(size_t index, Optional<IntSize>)
{
    if (index > 0)
        return Error::from_string_literal("ILBMImageDecoderPlugin: frame index must be 0");

    if (m_context->state < ILBMLoadingContext::State::BitmapDecoded)
        TRY(decode_iff_chunks(*m_context));

    VERIFY(m_context->bitmap);
    return ImageFrameDescriptor { m_context->bitmap, 0 };
}

}