#include "assets/image_metadata.h"

#include <cstddef>

namespace pp::assets {

namespace {

constexpr std::byte png_signature[] {
    std::byte { 0x89 },
    std::byte { 0x50 },
    std::byte { 0x4e },
    std::byte { 0x47 },
    std::byte { 0x0d },
    std::byte { 0x0a },
    std::byte { 0x1a },
    std::byte { 0x0a },
};

[[nodiscard]] bool has_png_signature(std::span<const std::byte> bytes) noexcept
{
    if (bytes.size() < 8U) {
        return false;
    }

    for (std::size_t i = 0; i < 8U; ++i) {
        if (bytes[i] != png_signature[i]) {
            return false;
        }
    }

    return true;
}

[[nodiscard]] std::uint32_t read_u32_be(std::span<const std::byte> bytes, std::size_t offset) noexcept
{
    return (static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(bytes[offset])) << 24U)
        | (static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(bytes[offset + 1U])) << 16U)
        | (static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(bytes[offset + 2U])) << 8U)
        | static_cast<std::uint32_t>(std::to_integer<std::uint8_t>(bytes[offset + 3U]));
}

[[nodiscard]] pp::foundation::Result<ImageColorType> parse_png_color_type(std::byte value) noexcept
{
    switch (std::to_integer<std::uint8_t>(value)) {
    case 0:
        return pp::foundation::Result<ImageColorType>::success(ImageColorType::grayscale);
    case 2:
        return pp::foundation::Result<ImageColorType>::success(ImageColorType::rgb);
    case 3:
        return pp::foundation::Result<ImageColorType>::success(ImageColorType::indexed);
    case 4:
        return pp::foundation::Result<ImageColorType>::success(ImageColorType::grayscale_alpha);
    case 6:
        return pp::foundation::Result<ImageColorType>::success(ImageColorType::rgba);
    default:
        return pp::foundation::Result<ImageColorType>::failure(
            pp::foundation::Status::invalid_argument("PNG color type is unsupported"));
    }
}

[[nodiscard]] std::uint8_t component_count(ImageColorType color_type) noexcept
{
    switch (color_type) {
    case ImageColorType::grayscale:
    case ImageColorType::indexed:
        return 1;
    case ImageColorType::grayscale_alpha:
        return 2;
    case ImageColorType::rgb:
        return 3;
    case ImageColorType::rgba:
        return 4;
    }

    return 0;
}

}

pp::foundation::Result<ImageMetadata> parse_png_metadata(std::span<const std::byte> bytes) noexcept
{
    constexpr std::size_t png_ihdr_end = 33;
    if (bytes.size() < png_ihdr_end) {
        return pp::foundation::Result<ImageMetadata>::failure(
            pp::foundation::Status::out_of_range("PNG metadata is truncated"));
    }

    if (!has_png_signature(bytes)) {
        return pp::foundation::Result<ImageMetadata>::failure(
            pp::foundation::Status::invalid_argument("PNG signature is invalid"));
    }

    const auto ihdr_length = read_u32_be(bytes, 8);
    if (ihdr_length != 13U || bytes[12] != std::byte { 'I' } || bytes[13] != std::byte { 'H' }
        || bytes[14] != std::byte { 'D' } || bytes[15] != std::byte { 'R' }) {
        return pp::foundation::Result<ImageMetadata>::failure(
            pp::foundation::Status::invalid_argument("PNG IHDR chunk is invalid"));
    }

    const auto width = read_u32_be(bytes, 16);
    const auto height = read_u32_be(bytes, 20);
    if (width == 0 || height == 0 || width > max_image_dimension || height > max_image_dimension) {
        return pp::foundation::Result<ImageMetadata>::failure(
            pp::foundation::Status::out_of_range("PNG dimensions are outside the configured range"));
    }

    const auto bit_depth = std::to_integer<std::uint8_t>(bytes[24]);
    if (bit_depth == 0U) {
        return pp::foundation::Result<ImageMetadata>::failure(
            pp::foundation::Status::invalid_argument("PNG bit depth is invalid"));
    }

    const auto color_type = parse_png_color_type(bytes[25]);
    if (!color_type) {
        return pp::foundation::Result<ImageMetadata>::failure(color_type.status());
    }

    return pp::foundation::Result<ImageMetadata>::success(ImageMetadata {
        .width = width,
        .height = height,
        .bit_depth = bit_depth,
        .components = component_count(color_type.value()),
        .color_type = color_type.value(),
    });
}

const char* image_color_type_name(ImageColorType color_type) noexcept
{
    switch (color_type) {
    case ImageColorType::grayscale:
        return "grayscale";
    case ImageColorType::rgb:
        return "rgb";
    case ImageColorType::indexed:
        return "indexed";
    case ImageColorType::grayscale_alpha:
        return "grayscale_alpha";
    case ImageColorType::rgba:
        return "rgba";
    }

    return "unknown";
}

}
