#include "Pixel_Frame.h"

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

namespace renderive::web {
namespace {

void write_u32_le(char* target, std::uint32_t value) {
    target[0] = static_cast<char>(value & 0xffU);
    target[1] = static_cast<char>((value >> 8U) & 0xffU);
    target[2] = static_cast<char>((value >> 16U) & 0xffU);
    target[3] = static_cast<char>((value >> 24U) & 0xffU);
}

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, 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);
    if (width > (std::numeric_limits<std::size_t>::max() - pixel_frame_header_size) /
                    (height * sizeof(Pixel))) {
        return {};
    }

    const std::size_t pixel_bytes = width * height * sizeof(Pixel);
    std::string frame(pixel_frame_header_size + pixel_bytes, '\0');
    frame[0] = 'R';
    frame[1] = 'V';
    frame[2] = 'P';
    frame[3] = '1';
    write_u32_le(frame.data() + 4, static_cast<std::uint32_t>(image.width));
    write_u32_le(frame.data() + 8, static_cast<std::uint32_t>(image.height));
    write_u32_le(frame.data() + 12,
                 static_cast<std::uint32_t>(image.width * static_cast<int>(sizeof(Pixel))));

    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
