#include "assets/image_pixels.h"

#include "assets/image_metadata.h"

#include <limits>
#include <utility>

#define STB_IMAGE_STATIC
#define STB_IMAGE_IMPLEMENTATION
#include <stb/stb_image.h>

#define STB_IMAGE_WRITE_STATIC
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb/stb_image_write.h>

namespace pp::assets {

namespace {

[[nodiscard]] pp::foundation::Result<std::size_t> rgba_byte_size(
    std::uint32_t width,
    std::uint32_t height) noexcept
{
    const auto pixels = static_cast<std::uint64_t>(width) * static_cast<std::uint64_t>(height);
    constexpr auto channels = 4ULL;
    if (pixels > std::numeric_limits<std::uint64_t>::max() / channels) {
        return pp::foundation::Result<std::size_t>::failure(
            pp::foundation::Status::out_of_range("RGBA byte size overflows"));
    }

    const auto bytes = pixels * channels;
    if (bytes > static_cast<std::uint64_t>(std::numeric_limits<std::size_t>::max())) {
        return pp::foundation::Result<std::size_t>::failure(
            pp::foundation::Status::out_of_range("RGBA byte size exceeds addressable memory"));
    }

    return pp::foundation::Result<std::size_t>::success(static_cast<std::size_t>(bytes));
}

void append_png_bytes(void* context, void* data, int size)
{
    if (context == nullptr || data == nullptr || size <= 0) {
        return;
    }

    auto* bytes = static_cast<std::vector<std::byte>*>(context);
    const auto* begin = static_cast<const std::byte*>(data);
    bytes->insert(bytes->end(), begin, begin + size);
}

}

pp::foundation::Result<Rgba8Image> decode_png_rgba8(std::span<const std::byte> bytes)
{
    const auto metadata = parse_png_metadata(bytes);
    if (!metadata) {
        return pp::foundation::Result<Rgba8Image>::failure(metadata.status());
    }

    if (bytes.size() > static_cast<std::size_t>(std::numeric_limits<int>::max())) {
        return pp::foundation::Result<Rgba8Image>::failure(
            pp::foundation::Status::out_of_range("PNG payload is too large for the decoder"));
    }

    int width = 0;
    int height = 0;
    int source_components = 0;
    auto* decoded = stbi_load_from_memory(
        reinterpret_cast<const stbi_uc*>(bytes.data()),
        static_cast<int>(bytes.size()),
        &width,
        &height,
        &source_components,
        4);
    if (decoded == nullptr) {
        return pp::foundation::Result<Rgba8Image>::failure(
            pp::foundation::Status::invalid_argument("PNG payload could not be decoded"));
    }

    const auto cleanup = [decoded]() noexcept {
        stbi_image_free(decoded);
    };

    if (width <= 0 || height <= 0
        || static_cast<std::uint32_t>(width) != metadata.value().width
        || static_cast<std::uint32_t>(height) != metadata.value().height) {
        cleanup();
        return pp::foundation::Result<Rgba8Image>::failure(
            pp::foundation::Status::invalid_argument("decoded PNG dimensions are inconsistent"));
    }

    const auto byte_count = rgba_byte_size(metadata.value().width, metadata.value().height);
    if (!byte_count) {
        cleanup();
        return pp::foundation::Result<Rgba8Image>::failure(byte_count.status());
    }

    Rgba8Image image {
        .width = metadata.value().width,
        .height = metadata.value().height,
        .pixels = {},
    };
    image.pixels.assign(decoded, decoded + byte_count.value());
    cleanup();

    return pp::foundation::Result<Rgba8Image>::success(std::move(image));
}

pp::foundation::Result<std::vector<std::byte>> encode_png_rgba8(
    std::uint32_t width,
    std::uint32_t height,
    std::span<const std::uint8_t> pixels)
{
    if (width == 0 || height == 0) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(
            pp::foundation::Status::invalid_argument("PNG dimensions must be greater than zero"));
    }

    if (width > static_cast<std::uint32_t>(std::numeric_limits<int>::max())
        || height > static_cast<std::uint32_t>(std::numeric_limits<int>::max())) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(
            pp::foundation::Status::out_of_range("PNG dimensions exceed encoder limits"));
    }

    const auto byte_count = rgba_byte_size(width, height);
    if (!byte_count) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(byte_count.status());
    }

    if (pixels.size() != byte_count.value()) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(
            pp::foundation::Status::invalid_argument("RGBA pixel payload size does not match dimensions"));
    }

    const auto stride = static_cast<std::uint64_t>(width) * 4ULL;
    if (stride > static_cast<std::uint64_t>(std::numeric_limits<int>::max())) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(
            pp::foundation::Status::out_of_range("PNG row stride exceeds encoder limits"));
    }

    std::vector<std::byte> encoded;
    const auto result = stbi_write_png_to_func(
        append_png_bytes,
        &encoded,
        static_cast<int>(width),
        static_cast<int>(height),
        4,
        pixels.data(),
        static_cast<int>(stride));

    if (result == 0 || encoded.empty()) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(
            pp::foundation::Status::invalid_argument("RGBA pixels could not be encoded as PNG"));
    }

    return pp::foundation::Result<std::vector<std::byte>>::success(std::move(encoded));
}

}