#include "NativeBlurHash.hpp"

#include <algorithm>
#include <array>
#include <cmath>

namespace {

constexpr char kBase83Chars[] =
    "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz#$%*+,-.:;=?@[]^_{|}~";
constexpr float kPi = 3.14159265358979323846f;

constexpr std::array<std::int8_t, 128> make_base83_lookup() {
    std::array<std::int8_t, 128> lookup = {};
    for (std::size_t i = 0; i < lookup.size(); i++) {
        lookup[i] = -1;
    }
    for (int i = 0; i < 83; i++) {
        lookup[static_cast<unsigned char>(kBase83Chars[i])] = static_cast<std::int8_t>(i);
    }
    return lookup;
}

constexpr auto kBase83Lookup = make_base83_lookup();

int base83_index(char c) {
    const auto uc = static_cast<unsigned char>(c);
    if (uc >= kBase83Lookup.size()) {
        return 0;
    }
    const int index = kBase83Lookup[uc];
    return index >= 0 ? index : 0;
}

int decode_base83(const std::string &str, std::size_t start, std::size_t length) {
    int value = 0;
    for (std::size_t i = 0; i < length; i++) {
        value = value * 83 + base83_index(str[start + i]);
    }
    return value;
}

float srgb_to_linear(int value) {
    const float v = static_cast<float>(value) / 255.0f;
    if (v <= 0.04045f) {
        return v / 12.92f;
    }
    return std::pow((v + 0.055f) / 1.055f, 2.4f);
}

int linear_to_srgb(float value) {
    const float v = std::clamp(value, 0.0f, 1.0f);
    if (v <= 0.0031308f) {
        return static_cast<int>(v * 12.92f * 255.0f + 0.5f);
    }
    return static_cast<int>((1.055f * std::pow(v, 1.0f / 2.4f) - 0.055f) * 255.0f + 0.5f);
}

float sign_pow(float value, float exp) {
    return std::copysign(std::pow(std::abs(value), exp), value);
}

std::array<float, 3> decode_dc(int value) {
    const int r = value >> 16;
    const int g = (value >> 8) & 255;
    const int b = value & 255;
    return {srgb_to_linear(r), srgb_to_linear(g), srgb_to_linear(b)};
}

std::array<float, 3> decode_ac(int value, float maximum_value) {
    const int quant_r = value / (19 * 19);
    const int quant_g = (value / 19) % 19;
    const int quant_b = value % 19;
    return {
        sign_pow((static_cast<float>(quant_r) - 9.0f) / 9.0f, 2.0f) * maximum_value,
        sign_pow((static_cast<float>(quant_g) - 9.0f) / 9.0f, 2.0f) * maximum_value,
        sign_pow((static_cast<float>(quant_b) - 9.0f) / 9.0f, 2.0f) * maximum_value
    };
}

} // namespace

namespace rosetta::nativeimage {

bool decode_blurhash_to_rgb(
    const std::string &blurhash,
    int width,
    int height,
    float punch,
    std::vector<std::uint8_t> &rgb
) {
    if (blurhash.size() < 6 || width <= 0 || height <= 0) {
        return false;
    }

    const int size_flag = decode_base83(blurhash, 0, 1);
    const int num_y = (size_flag / 9) + 1;
    const int num_x = (size_flag % 9) + 1;
    const std::size_t expected_length = static_cast<std::size_t>(4 + 2 * num_x * num_y);
    if (blurhash.size() != expected_length) {
        return false;
    }

    const int quantized_maximum = decode_base83(blurhash, 1, 1);
    const float maximum_value = static_cast<float>(quantized_maximum + 1) / 166.0f;

    std::vector<std::array<float, 3>> colors;
    colors.reserve(static_cast<std::size_t>(num_x * num_y));
    for (int i = 0; i < num_x * num_y; i++) {
        if (i == 0) {
            colors.push_back(decode_dc(decode_base83(blurhash, 2, 4)));
        } else {
            colors.push_back(decode_ac(
                decode_base83(blurhash, static_cast<std::size_t>(4 + i * 2), 2),
                maximum_value * punch
            ));
        }
    }

    const std::size_t pixel_count = static_cast<std::size_t>(width) * static_cast<std::size_t>(height);
    rgb.assign(pixel_count * 3, 0);

    std::vector<float> cos_x;
    std::vector<float> cos_y;
    cos_x.resize(static_cast<std::size_t>(width) * static_cast<std::size_t>(num_x));
    cos_y.resize(static_cast<std::size_t>(height) * static_cast<std::size_t>(num_y));

    for (int x = 0; x < width; x++) {
        for (int i = 0; i < num_x; i++) {
            cos_x[static_cast<std::size_t>(x * num_x + i)] = std::cos(
                kPi * static_cast<float>(x) * static_cast<float>(i)
                / static_cast<float>(width)
            );
        }
    }
    for (int y = 0; y < height; y++) {
        for (int j = 0; j < num_y; j++) {
            cos_y[static_cast<std::size_t>(y * num_y + j)] = std::cos(
                kPi * static_cast<float>(y) * static_cast<float>(j)
                / static_cast<float>(height)
            );
        }
    }

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            float r = 0;
            float g = 0;
            float b = 0;

            for (int j = 0; j < num_y; j++) {
                const float basis_y = cos_y[static_cast<std::size_t>(y * num_y + j)];
                for (int i = 0; i < num_x; i++) {
                    const float basis_x = cos_x[static_cast<std::size_t>(x * num_x + i)];
                    const float basis = basis_x * basis_y;
                    const auto &color = colors[static_cast<std::size_t>(i + j * num_x)];
                    r += color[0] * basis;
                    g += color[1] * basis;
                    b += color[2] * basis;
                }
            }

            const std::size_t offset = static_cast<std::size_t>((y * width + x) * 3);
            rgb[offset] = static_cast<std::uint8_t>(std::clamp(linear_to_srgb(r), 0, 255));
            rgb[offset + 1] = static_cast<std::uint8_t>(std::clamp(linear_to_srgb(g), 0, 255));
            rgb[offset + 2] = static_cast<std::uint8_t>(std::clamp(linear_to_srgb(b), 0, 255));
        }
    }

    return true;
}

} // namespace rosetta::nativeimage
