#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);
}

void append_encoded_bytes(void* context, void* data, int size)
{
    append_png_bytes(context, data, size);
}

[[nodiscard]] bool has_jpeg_soi(std::span<const std::byte> bytes) noexcept
{
    return bytes.size() >= 2U
        && bytes[0] == std::byte { 0xff }
        && bytes[1] == std::byte { 0xd8 };
}

constexpr char gpano_xmp[] =
    "http://ns.adobe.com/xap/1.0/\0"
    R"(<?xpacket begin="" id="W5M0MpCehiHzreSzNTczkc9d"?>
  <x:xmpmeta xmlns:x="adobe:ns:meta/" xmptk="SAMSUNG 360CAM">
    <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">
      <rdf:Description rdf:about="" xmlns:GPano="http://ns.google.com/photos/1.0/panorama/">
        <GPano:ProjectionType>equirectangular</GPano:ProjectionType>
        <GPano:UsePanoramaViewer>True</GPano:UsePanoramaViewer>
        <GPano:CroppedAreaLeftPixels>0</GPano:CroppedAreaLeftPixels>
        <GPano:CroppedAreaTopPixels>0</GPano:CroppedAreaTopPixels>
        <GPano:PoseHeadingDegrees>0</GPano:PoseHeadingDegrees>
        <GPano:PosePitchDegrees>0</GPano:PosePitchDegrees>
        <GPano:PoseRollDegrees>0</GPano:PoseRollDegrees>
        <GPano:StitchingSoftware>PanoPainter</GPano:StitchingSoftware>
      </rdf:Description>
    </rdf:RDF>
  </x:xmpmeta>
<?xpacket end="r"?>)";

}

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<Rgba8Image> decode_jpeg_rgba8(std::span<const std::byte> bytes)
{
    if (!has_jpeg_soi(bytes)) {
        return pp::foundation::Result<Rgba8Image>::failure(
            pp::foundation::Status::invalid_argument("JPEG signature is invalid"));
    }

    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("JPEG 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("JPEG 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) > max_image_dimension
        || static_cast<std::uint32_t>(height) > max_image_dimension) {
        cleanup();
        return pp::foundation::Result<Rgba8Image>::failure(
            pp::foundation::Status::out_of_range("decoded JPEG dimensions are outside the configured range"));
    }

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

    Rgba8Image image {
        .width = static_cast<std::uint32_t>(width),
        .height = static_cast<std::uint32_t>(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));
}

pp::foundation::Result<std::vector<std::byte>> encode_jpeg_rgba8(
    std::uint32_t width,
    std::uint32_t height,
    std::span<const std::uint8_t> pixels,
    int quality)
{
    if (width == 0 || height == 0) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(
            pp::foundation::Status::invalid_argument("JPEG 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("JPEG dimensions exceed encoder limits"));
    }

    if (quality < 1 || quality > 100) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(
            pp::foundation::Status::out_of_range("JPEG quality must be within 1..100"));
    }

    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"));
    }

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

    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 JPEG"));
    }

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

pp::foundation::Result<std::vector<std::byte>> inject_gpano_xmp_into_jpeg(std::span<const std::byte> jpeg)
{
    if (!has_jpeg_soi(jpeg)) {
        return pp::foundation::Result<std::vector<std::byte>>::failure(
            pp::foundation::Status::invalid_argument("JPEG signature is invalid"));
    }

    constexpr auto xmp_size = sizeof(gpano_xmp);
    static_assert(xmp_size + 2U <= 0xffffU);
    const auto segment_length = static_cast<std::uint16_t>(xmp_size + 2U);

    std::vector<std::byte> encoded;
    encoded.reserve(jpeg.size() + xmp_size + 4U);
    encoded.insert(encoded.end(), jpeg.begin(), jpeg.begin() + 2);
    encoded.push_back(std::byte { 0xff });
    encoded.push_back(std::byte { 0xe1 });
    encoded.push_back(static_cast<std::byte>((segment_length >> 8U) & 0xffU));
    encoded.push_back(static_cast<std::byte>(segment_length & 0xffU));
    const auto* xmp_begin = reinterpret_cast<const std::byte*>(gpano_xmp);
    encoded.insert(encoded.end(), xmp_begin, xmp_begin + xmp_size);
    encoded.insert(encoded.end(), jpeg.begin() + 2, jpeg.end());

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

}
