#include "Pixel_Frame.h"

#include <algorithm>
#if defined(_M_X64) || defined(__x86_64__)
#include <tmmintrin.h>
#endif

namespace renderive::web {
namespace {

std::uint8_t flatten(std::uint8_t premultiplied, std::uint8_t alpha,
                     std::uint8_t background) {
    return static_cast<std::uint8_t>(
        std::min(255U, static_cast<unsigned>(premultiplied) +
                           (static_cast<unsigned>(background) * (255U - alpha) + 127U) / 255U));
}

} // namespace

std::string encode_pixel_frame(Image_View image, std::uint64_t sequence, Color background) {
    if (image.empty() || image.format != Pixel_Format::Premultiplied_32 ||
        image.stride < image.width * static_cast<int>(sizeof(Pixel))) {
        return {};
    }
    const auto width = static_cast<std::size_t>(image.width);
    const auto height = static_cast<std::size_t>(image.height);
    std::string frame = make_rgba8_pixel_frame(
        static_cast<std::uint32_t>(width), static_cast<std::uint32_t>(height), sequence);
    if (frame.empty())
        return {};
    auto* output = reinterpret_cast<std::uint32_t*>(frame.data() + pixel_frame_header_size);
    const std::uint32_t opaque_background =
        static_cast<std::uint32_t>(background.r) |
        (static_cast<std::uint32_t>(background.g) << 8U) |
        (static_cast<std::uint32_t>(background.b) << 16U) |
        0xff000000U;
    const auto convert_pixel = [background, opaque_background](Pixel pixel) {
        const auto alpha = static_cast<std::uint8_t>((pixel >> 24U) & 0xffU);
        if (alpha == 0)
            return opaque_background;
        const auto red = static_cast<std::uint8_t>((pixel >> 16U) & 0xffU);
        const auto green = static_cast<std::uint8_t>((pixel >> 8U) & 0xffU);
        const auto blue = static_cast<std::uint8_t>(pixel & 0xffU);
        if (alpha == 255) {
            return static_cast<std::uint32_t>(red) |
                   (static_cast<std::uint32_t>(green) << 8U) |
                   (static_cast<std::uint32_t>(blue) << 16U) |
                   0xff000000U;
        }
        return static_cast<std::uint32_t>(flatten(red, alpha, background.r)) |
               (static_cast<std::uint32_t>(flatten(green, alpha, background.g)) << 8U) |
               (static_cast<std::uint32_t>(flatten(blue, alpha, background.b)) << 16U) |
               0xff000000U;
    };
#if defined(_M_X64) || defined(__x86_64__)
    const __m128i rgba_shuffle = _mm_setr_epi8(
        2, 1, 0, 3, 6, 5, 4, 7, 10, 9, 8, 11, 14, 13, 12, 15);
    const __m128i background_pixels = _mm_set1_epi32(static_cast<int>(opaque_background));
#endif
    for (int y = 0; y < image.height; ++y) {
        const auto* row = reinterpret_cast<const Pixel*>(
            image.data + static_cast<std::ptrdiff_t>(y) * image.stride);
        int x{};
#if defined(_M_X64) || defined(__x86_64__)
        for (; x + 4 <= image.width; x += 4) {
            const Pixel alpha_union = row[x] | row[x + 1] | row[x + 2] | row[x + 3];
            if ((alpha_union & 0xff000000U) == 0) {
                _mm_storeu_si128(reinterpret_cast<__m128i*>(output), background_pixels);
                output += 4;
                continue;
            }
            const Pixel alpha_intersection = row[x] & row[x + 1] & row[x + 2] & row[x + 3];
            if ((alpha_intersection & 0xff000000U) == 0xff000000U) {
                const __m128i bgra =
                    _mm_loadu_si128(reinterpret_cast<const __m128i*>(row + x));
                _mm_storeu_si128(reinterpret_cast<__m128i*>(output),
                                 _mm_shuffle_epi8(bgra, rgba_shuffle));
                output += 4;
                continue;
            }
            *output++ = convert_pixel(row[x]);
            *output++ = convert_pixel(row[x + 1]);
            *output++ = convert_pixel(row[x + 2]);
            *output++ = convert_pixel(row[x + 3]);
        }
#endif
        for (; x < image.width; ++x)
            *output++ = convert_pixel(row[x]);
    }
    return frame;
}

} // namespace renderive::web