/*
 * Copyright 2023 Rive
 */

#include "rive/decoders/bitmap_decoder.hpp"
#include "rive/rive_types.hpp"
#include "rive/math/simd.hpp"
#include <stdio.h>
#include <string.h>

Bitmap::Bitmap(uint32_t width,
               uint32_t height,
               size_t numBytes,
               PixelFormat pixelFormat,
               std::unique_ptr<const uint8_t[]> bytes) :
    m_Width(width),
    m_Height(height),
    m_NumBytes(numBytes),
    m_PixelFormat(pixelFormat),
    m_Bytes(std::move(bytes))
{}

static size_t bytes_per_pixel(Bitmap::PixelFormat format)
{
    switch (format)
    {
        case Bitmap::PixelFormat::RGB:
            return 3;
        case Bitmap::PixelFormat::RGBA:
        case Bitmap::PixelFormat::RGBAPremul:
            return 4;
    }
    RIVE_UNREACHABLE();
}

void Bitmap::pixelFormat(PixelFormat format)
{
    if (format == m_PixelFormat)
    {
        return;
    }

    size_t imageNumPixels = m_Height * m_Width;
    size_t fromBytesPerPixel = bytes_per_pixel(m_PixelFormat);
    size_t toBytesPerPixel = bytes_per_pixel(format);
    size_t toSizeInBytes = imageNumPixels * toBytesPerPixel;

    // Round our allocation up to a multiple of 8 so we can premultiply two
    // pixels at once if needed.
    size_t allocSize = (toSizeInBytes + 7) & ~7llu;
    assert(allocSize >= toSizeInBytes && allocSize <= toSizeInBytes + 7);
    auto toBytes = std::unique_ptr<uint8_t[]>(new uint8_t[allocSize]);

    // Copy into the new pixel buffer.
    int writeIndex = 0;
    int readIndex = 0;
    for (size_t i = 0; i < imageNumPixels; i++)
    {
        for (size_t j = 0; j < toBytesPerPixel; j++)
        {
            toBytes[writeIndex++] =
                j < fromBytesPerPixel ? m_Bytes[readIndex++] : 255;
        }
    }

    if (format == PixelFormat::RGBAPremul)
    {
        // Convert to premultiplied alpha.
        for (size_t i = 0; i < toSizeInBytes; i += 8)
        {
            // Load 2 pixels into 64 bits. We overallocated our buffer to ensure
            // we could always premultiply in twos.
            assert(i + 8 <= allocSize);
            rive::uint8x8 twoPx = rive::simd::load<uint8_t, 8>(&toBytes[i]);
            uint8_t a0 = twoPx[3];
            uint8_t a1 = twoPx[7];
            // Don't premultiply fully opaque pixels.
            if (a0 != 255 || a1 != 255)
            {
                // Cast to 16 bits to avoid overflow.
                rive::uint16x8 widePx = rive::simd::cast<uint16_t>(twoPx);
                // Multiply by alpha.
                rive::uint16x8 alpha = {a0, a0, a0, 255, a1, a1, a1, 255};
                widePx = rive::simd::div255(widePx * alpha);
                // Cast back to 8 bits and store.
                twoPx = rive::simd::cast<uint8_t>(widePx);
                rive::simd::store(&toBytes[i], twoPx);
            }
        }
    }
    else
    {
        // Unmultiplying alpha is not currently supported.
        assert(m_PixelFormat != PixelFormat::RGBAPremul);
    }

    m_Bytes = std::move(toBytes);
    m_PixelFormat = format;
    m_NumBytes = toSizeInBytes;
}