/*

 * Copyright (c) 2020, Paul Roukema <roukemap@gmail.com>

 *

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

 */

#include <AK/ByteBuffer.h>

#include <AK/Debug.h>

#include <AK/Enumerate.h>

#include <AK/MemoryStream.h>

#include <AK/Types.h>

#include <LibGfx/ImageFormats/BMPLoader.h>

#include <LibGfx/ImageFormats/ICOLoader.h>

#include <LibGfx/ImageFormats/PNGLoader.h>

#include <string.h>

namespace Gfx {

// FIXME: This is in little-endian order. Maybe need a NetworkOrdered<T> equivalent eventually.

struct ICONDIR {

    u16 must_be_0 = 0;

    u16 must_be_1 = 0;

    u16 image_count = 0;

};

static_assert(AssertSize<ICONDIR, 6>());

struct ICONDIRENTRY {

    u8 width;

    u8 height;

    u8 color_count;

    u8 reserved_0;

    u16 planes;

    u16 bits_per_pixel;

    u32 size;

    u32 offset;

};

static_assert(AssertSize<ICONDIRENTRY, 16>());

};

template<>

class AK::Traits<Gfx::ICONDIR> : public DefaultTraits<Gfx::ICONDIR> {

public:

    static constexpr bool is_trivially_serializable() { return true; }

};

template<>

class AK::Traits<Gfx::ICONDIRENTRY> : public DefaultTraits<Gfx::ICONDIRENTRY> {

public:

    static constexpr bool is_trivially_serializable() { return true; }

};

namespace Gfx {

struct ICOImageDescriptor {

    u16 width;

    u16 height;

    u16 bits_per_pixel;

    size_t offset;

    size_t size;

    RefPtr<Gfx::Bitmap> bitmap;

};

struct ICOLoadingContext {

    enum State {

        NotDecoded = 0,

        Error,

        DirectoryDecoded,

        BitmapDecoded

    };

    State state { NotDecoded };

    u8 const* data { nullptr };

    size_t data_size { 0 };

    Vector<ICOImageDescriptor> images;

    size_t largest_index;

};

static ErrorOr<size_t> decode_ico_header(Stream& stream)

{

    auto header = TRY(stream.read_value<ICONDIR>());

    if (header.must_be_0 != 0 || header.must_be_1 != 1)

        return Error::from_string_literal("Invalid ICO header");

    return { header.image_count };

}

static ErrorOr<ICOImageDescriptor> decode_ico_direntry(Stream& stream)

{

    auto entry = TRY(stream.read_value<ICONDIRENTRY>());

    ICOImageDescriptor desc = { entry.width, entry.height, entry.bits_per_pixel, entry.offset, entry.size, nullptr };

    if (desc.width == 0)

        desc.width = 256;

    if (desc.height == 0)

        desc.height = 256;

    return { desc };

}

static size_t find_largest_image(ICOLoadingContext const& context)

{

    size_t max_area = 0;

    size_t largest_index = 0;

    u16 max_bits_per_pixel = 0;

    for (auto const& [index, desc] : enumerate(context.images)) {

        auto area = static_cast<size_t>(desc.width) * static_cast<size_t>(desc.height);

        if (area > max_area || (area == max_area && desc.bits_per_pixel > max_bits_per_pixel)) {

            max_area = area;

            largest_index = index;

            max_bits_per_pixel = desc.bits_per_pixel;

        }

    }

    return largest_index;

}

static ErrorOr<void> load_ico_directory(ICOLoadingContext& context)

{

    FixedMemoryStream stream { { context.data, context.data_size } };

    auto image_count = TRY(decode_ico_header(stream));

    if (image_count == 0)

        return Error::from_string_literal("ICO file has no images");

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

        auto desc = TRY(decode_ico_direntry(stream));

        if (desc.offset + desc.size < desc.offset // detect integer overflow

            || (desc.offset + desc.size) > context.data_size) {

            dbgln_if(ICO_DEBUG, "load_ico_directory: offset: {} size: {} doesn't fit in ICO size: {}", desc.offset, desc.size, context.data_size);

            return Error::from_string_literal("ICO size too large");

        }

        dbgln_if(ICO_DEBUG, "load_ico_directory: index {} width: {} height: {} offset: {} size: {}", i, desc.width, desc.height, desc.offset, desc.size);

        TRY(context.images.try_append(desc));

    }

    context.largest_index = find_largest_image(context);

    context.state = ICOLoadingContext::State::DirectoryDecoded;

    return {};

}

ErrorOr<void> ICOImageDecoderPlugin::load_ico_bitmap(ICOLoadingContext& context)

{

    VERIFY(context.state >= ICOLoadingContext::State::DirectoryDecoded);

    size_t const real_index = context.largest_index;

    if (real_index >= context.images.size())

        return Error::from_string_literal("Index out of bounds");

    ICOImageDescriptor& desc = context.images[real_index];

    if (PNGImageDecoderPlugin::sniff({ context.data + desc.offset, desc.size })) {

        auto png_decoder = TRY(PNGImageDecoderPlugin::create({ context.data + desc.offset, desc.size }));

        auto decoded_png_frame = TRY(png_decoder->frame(0));

        if (!decoded_png_frame.image) {

            dbgln_if(ICO_DEBUG, "load_ico_bitmap: failed to load PNG encoded image index: {}", real_index);

            return Error::from_string_literal("Encoded image not null");

        }

        desc.bitmap = decoded_png_frame.image;

        return {};

    } else {

        auto bmp_decoder = TRY(BMPImageDecoderPlugin::create_as_included_in_ico({}, { context.data + desc.offset, desc.size }));

        // NOTE: We don't initialize a BMP decoder in the usual way, but rather

        // we just create an object and try to sniff for a frame when it's included

        // inside an ICO image.

        if (bmp_decoder->sniff_dib()) {

            auto decoded_bmp_frame = TRY(bmp_decoder->frame(0));

            if (!decoded_bmp_frame.image) {

                dbgln_if(ICO_DEBUG, "load_ico_bitmap: failed to load BMP encoded image index: {}", real_index);

                return Error::from_string_literal("Encoded image not null");

            }

            desc.bitmap = decoded_bmp_frame.image;

        } else {

            dbgln_if(ICO_DEBUG, "load_ico_bitmap: encoded image not supported at index: {}", real_index);

            return Error::from_string_literal("Encoded image not supported");

        }

        return {};

    }

}

bool ICOImageDecoderPlugin::sniff(ReadonlyBytes data)

{

    FixedMemoryStream stream { data };

    return !decode_ico_header(stream).is_error();

}

ErrorOr<NonnullOwnPtr<ImageDecoderPlugin>> ICOImageDecoderPlugin::create(ReadonlyBytes data)

{

    auto plugin = TRY(adopt_nonnull_own_or_enomem(new (nothrow) ICOImageDecoderPlugin(data.data(), data.size())));

    TRY(load_ico_directory(*plugin->m_context));

    return plugin;

}

ICOImageDecoderPlugin::ICOImageDecoderPlugin(u8 const* data, size_t size)

{

    m_context = make<ICOLoadingContext>();

    m_context->data = data;

    m_context->data_size = size;

}

ICOImageDecoderPlugin::~ICOImageDecoderPlugin() = default;

IntSize ICOImageDecoderPlugin::size()

{

    return { m_context->images[m_context->largest_index].width, m_context->images[m_context->largest_index].height };

}

ErrorOr<ImageFrameDescriptor> ICOImageDecoderPlugin::frame(size_t index, Optional<IntSize>)

{

    if (index > 0)

        return Error::from_string_literal("ICOImageDecoderPlugin: Invalid frame index");

    if (m_context->state == ICOLoadingContext::State::Error)

        return Error::from_string_literal("ICOImageDecoderPlugin: Decoding failed");

    if (m_context->state < ICOLoadingContext::State::BitmapDecoded) {

        // NOTE: This forces the chunk decoding to happen.

        auto maybe_error = load_ico_bitmap(*m_context);

        if (maybe_error.is_error()) {

            m_context->state = ICOLoadingContext::State::Error;

            return Error::from_string_literal("ICOImageDecoderPlugin: Decoding failed");

        }

        m_context->state = ICOLoadingContext::State::BitmapDecoded;

    }

    VERIFY(m_context->images[m_context->largest_index].bitmap);

    return ImageFrameDescriptor { m_context->images[m_context->largest_index].bitmap, 0 };

}

}