/*

 * Copyright (c) 2022, Tom Needham <06needhamt@gmail.com>

 *

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

 */

#include <AK/MemoryStream.h>

#include <AK/Span.h>

#include <AK/StdLibExtraDetails.h>

#include <AK/String.h>

#include <AK/Traits.h>

#include <LibGfx/ImageFormats/TGALoader.h>

namespace Gfx {

enum TGADataType : u8 {

    None = 0,

    UncompressedColorMapped = 1,

    UncompressedRGB = 2,

    UncompressedBlackAndWhite = 3,

    RunLengthEncodedColorMapped = 9,

    RunLengthEncodedRGB = 10,

    CompressedBlackAndWhite = 11,

    CompressedColorMapped = 32,

    CompressedColorMappedFourPass = 33

};

struct [[gnu::packed]] TGAHeader {

    u8 id_length;

    u8 color_map_type;

    TGADataType data_type_code;

    i16 color_map_origin;

    i16 color_map_length;

    u8 color_map_depth;

    i16 x_origin;

    i16 y_origin;

    u16 width;

    u16 height;

    u8 bits_per_pixel;

    u8 image_descriptor;

};

static_assert(sizeof(TGAHeader) == 18);

}

class TGAImageIdentifier : public Gfx::Metadata {

public:

    TGAImageIdentifier(String identifier)

        : m_identifier(move(identifier))

    {

    }

protected:

    virtual void fill_main_tags() const override

    {

        if (!m_identifier.is_empty())

            m_main_tags.set("Identifier"sv, m_identifier);

    }

    String m_identifier;

};

template<>

struct AK::Traits<Gfx::TGAHeader> : public DefaultTraits<Gfx::TGAHeader> {

    static constexpr bool is_trivially_serializable() { return true; }

};

namespace Gfx {

struct TGALoadingContext {

    TGALoadingContext(ReadonlyBytes bytes, FixedMemoryStream stream)

        : bytes(bytes)

        , stream(move(stream))

    {

    }

    ReadonlyBytes bytes;

    FixedMemoryStream stream;

    TGAHeader header {};

    RefPtr<Gfx::Bitmap> bitmap;

    OwnPtr<TGAImageIdentifier> identifier_metadata;

};

TGAImageDecoderPlugin::TGAImageDecoderPlugin(NonnullOwnPtr<TGALoadingContext> context)

    : m_context(move(context))

{

}

TGAImageDecoderPlugin::~TGAImageDecoderPlugin() = default;

IntSize TGAImageDecoderPlugin::size()

{

    return IntSize { m_context->header.width, m_context->header.height };

}

static ErrorOr<void> ensure_header_validity(TGAHeader const& header, size_t whole_image_stream_size)

{

    if ((header.bits_per_pixel % 8) != 0 || header.bits_per_pixel < 8 || header.bits_per_pixel > 32)

        return Error::from_string_literal("Invalid bit depth");

    auto bytes_remaining = whole_image_stream_size - sizeof(TGAHeader) - header.id_length;

    if (header.data_type_code == TGADataType::UncompressedRGB && bytes_remaining < static_cast<u64>(header.width) * header.height * (header.bits_per_pixel / 8))

        return Error::from_string_literal("Not enough data to read an image with the expected size");

    return {};

}

ErrorOr<void> TGAImageDecoderPlugin::decode_tga_header()

{

    m_context->header = TRY(m_context->stream.read_value<TGAHeader>());

    if (m_context->header.id_length > 0) {

        auto identifier = String::from_stream(m_context->stream, m_context->header.id_length);

        if (!identifier.is_error())

            m_context->identifier_metadata = make<TGAImageIdentifier>(identifier.release_value());

        else

            dbgln("FIXME: Handle non-utf8 identifiers in TGALoader");

    }

    TRY(ensure_header_validity(m_context->header, m_context->bytes.size()));

    return {};

}

bool TGAImageDecoderPlugin::validate_before_create(ReadonlyBytes data)

{

    FixedMemoryStream stream { data };

    auto header_or_err = stream.read_value<Gfx::TGAHeader>();

    if (header_or_err.is_error())

        return false;

    return !ensure_header_validity(header_or_err.release_value(), data.size()).is_error();

}

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

{

    FixedMemoryStream stream { data };

    auto context = TRY(adopt_nonnull_own_or_enomem(new (nothrow) TGALoadingContext(data, move(stream))));

    auto plugin = TRY(adopt_nonnull_own_or_enomem(new (nothrow) TGAImageDecoderPlugin(move(context))));

    TRY(plugin->decode_tga_header());

    return plugin;

}

static ErrorOr<ARGB32> read_pixel_from_stream(Stream& stream, size_t bytes_size)

{

    // NOTE: We support 8-bit, 24-bit and 32-bit color pixels

    VERIFY(bytes_size == 1 || bytes_size == 3 || bytes_size == 4);

    if (bytes_size == 1) {

        auto raw = TRY(stream.read_value<u8>());

        return Color(raw, raw, raw).value();

    }

    if (bytes_size == 3) {

        Array<u8, 3> raw;

        TRY(stream.read_until_filled(raw.span()));

        return Color(raw[2], raw[1], raw[0]).value();

    }

    return stream.read_value<ARGB32>();

}

struct TGAPixelPacketHeader {

    bool raw { false };

    u8 pixels_count { 0 };

};

static ErrorOr<TGAPixelPacketHeader> read_pixel_packet_header(Stream& stream)

{

    auto const pixel_packet_header = TRY(stream.read_value<u8>());

    bool pixels_raw_in_packet = !(pixel_packet_header & 0x80);

    u8 pixels_count_in_packet = (pixel_packet_header & 0x7f);

    // NOTE: Run-length-encoded/Raw pixel packets cannot encode zero pixels,

    // so value 0 stands for 1 pixel, 1 stands for 2, etc...

    pixels_count_in_packet++;

    VERIFY(pixels_count_in_packet > 0);

    return TGAPixelPacketHeader { pixels_raw_in_packet, pixels_count_in_packet };

}

Optional<Metadata const&> TGAImageDecoderPlugin::metadata()

{

    if (m_context->identifier_metadata)

        return *m_context->identifier_metadata;

    return OptionalNone {};

}

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

{

    auto bits_per_pixel = m_context->header.bits_per_pixel;

    auto color_map = m_context->header.color_map_type;

    auto data_type = m_context->header.data_type_code;

    auto width = m_context->header.width;

    auto height = m_context->header.height;

    auto image_descriptior = m_context->header.image_descriptor;

    if (index != 0)

        return Error::from_string_literal("TGAImageDecoderPlugin: frame index must be 0");

    if (color_map > 1)

        return Error::from_string_literal("TGAImageDecoderPlugin: Invalid color map type");

    if (m_context->bitmap)

        return ImageFrameDescriptor { m_context->bitmap, 0 };

    RefPtr<Gfx::Bitmap> bitmap;

    switch (bits_per_pixel) {

    case 8:

    case 24:

        bitmap = TRY(Bitmap::create(BitmapFormat::BGRx8888, { m_context->header.width, m_context->header.height }));

        break;

    case 32:

        bitmap = TRY(Bitmap::create(BitmapFormat::BGRA8888, { m_context->header.width, m_context->header.height }));

        break;

    default:

        // FIXME: Implement other TGA bit depths

        return Error::from_string_literal("TGAImageDecoderPlugin: Can only handle 8, 24 and 32 bits per pixel");

    }

    auto is_top_to_bottom = (image_descriptior & 1 << 5) == 0;

    auto is_left_to_right = (image_descriptior & 1 << 4) == 0;

    VERIFY((bits_per_pixel % 8) == 0);

    auto bytes_per_pixel = bits_per_pixel / 8;

    switch (data_type) {

    case TGADataType::UncompressedBlackAndWhite:

    case TGADataType::UncompressedRGB: {

        for (int row = 0; row < height; ++row) {

            for (int col = 0; col < width; ++col) {

                auto actual_row = row;

                if (is_top_to_bottom)

                    actual_row = height - 1 - row;

                auto actual_col = col;

                if (!is_left_to_right)

                    actual_col = width - 1 - col;

                bitmap->scanline(actual_row)[actual_col] = TRY(read_pixel_from_stream(m_context->stream, bytes_per_pixel));

            }

        }

        break;

    }

    case TGADataType::RunLengthEncodedRGB: {

        size_t pixel_index = 0;

        size_t pixel_count = height * width;

        while (pixel_index < pixel_count) {

            auto pixel_packet_header = TRY(read_pixel_packet_header(m_context->stream));

            VERIFY(pixel_packet_header.pixels_count > 0);

            auto pixel = TRY(read_pixel_from_stream(m_context->stream, bytes_per_pixel));

            auto max_pixel_index = min(pixel_index + pixel_packet_header.pixels_count, pixel_count);

            for (size_t current_pixel_index = pixel_index; current_pixel_index < max_pixel_index; ++current_pixel_index) {

                int row = current_pixel_index / width;

                int col = current_pixel_index % width;

                auto actual_row = row;

                if (is_top_to_bottom)

                    actual_row = height - 1 - row;

                auto actual_col = col;

                if (!is_left_to_right)

                    actual_col = width - 1 - col;

                bitmap->scanline(actual_row)[actual_col] = pixel;

                if (pixel_packet_header.raw && (current_pixel_index + 1) < max_pixel_index)

                    pixel = TRY(read_pixel_from_stream(m_context->stream, bytes_per_pixel));

            }

            pixel_index += pixel_packet_header.pixels_count;

        }

        break;

    }

    default:

        // FIXME: Implement other TGA data types

        return Error::from_string_literal("TGAImageDecoderPlugin: Can currently only handle the UncompressedRGB, CompressedRGB or UncompressedBlackAndWhite data type");

    }

    m_context->bitmap = bitmap;

    return ImageFrameDescriptor { m_context->bitmap, 0 };

}

}