// SPDX-FileCopyrightText: 2024 Joshua Goins <josh@redstrate.com>
// SPDX-License-Identifier: MIT

#include "blurhash.h"

#include <QColorSpace>

// From https://github.com/woltapp/blurhash/blob/master/Algorithm.md#base-83
const static QString b83Characters{QStringLiteral("0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz#$%*+,-.:;=?@[]^_{|}~")};

const static auto toLinearSRGB = QColorSpace(QColorSpace::SRgb).transformationToColorSpace(QColorSpace::SRgbLinear);
const static auto fromLinearSRGB = QColorSpace(QColorSpace::SRgbLinear).transformationToColorSpace(QColorSpace::SRgb);

QImage BlurHash::decode(const QString &blurhash, const QSize &size)
{
    // 10 is the minimum length of a blurhash string
    if (blurhash.length() < 10) {
        return {};
    }

    // First character is the number of components
    const auto components83 = decode83(blurhash.first(1));
    if (!components83.has_value()) {
        return {};
    }

    const auto components = unpackComponents(*components83);
    const auto minimumSize = 1 + 1 + 4 + (components.x * components.y - 1) * 2;
    if (components.x < 1 || components.y < 1 || blurhash.size() != minimumSize) {
        return {};
    }

    // Second character is the maximum AC component value
    const auto maxAC83 = decode83(blurhash.mid(1, 1));
    if (!maxAC83.has_value()) {
        return {};
    }

    const auto maxAC = decodeMaxAC(*maxAC83);

    // Third character onward is the average color of the image
    const auto averageColor83 = decode83(blurhash.mid(2, 4));
    if (!averageColor83.has_value()) {
        return {};
    }

    const auto averageColor = toLinearSRGB.map(decodeAverageColor(*averageColor83));

    std::vector<ColorF> values = {ColorF{.r = averageColor.redF(), .g = averageColor.greenF(), .b = averageColor.blueF()}};
    values.reserve(blurhash.size() - 7);

    // Iterate through the rest of the string for the color values
    // Each AC component is two characters each
    for (qsizetype c = 6; c < blurhash.size(); c += 2) {
        const auto acComponent83 = decode83(blurhash.mid(c, 2));
        if (!acComponent83.has_value()) {
            return {};
        }

        values.push_back(decodeAC(*acComponent83, maxAC));
    }

    QImage image(size, QImage::Format_RGB888);
    image.setColorSpace(QColorSpace::SRgb);

    const auto basisX = calculateWeights(size.width(), components.x);
    const auto basisY = calculateWeights(size.height(), components.y);

    for (int y = 0; y < size.height(); y++) {
        for (int x = 0; x < size.width(); x++) {
            float linearSumR = 0.0f;
            float linearSumG = 0.0f;
            float linearSumB = 0.0f;

            for (int nx = 0; nx < components.x; nx++) {
                for (int ny = 0; ny < components.y; ny++) {
                    const float basis = basisX.value(x * components.x + nx) * basisY.value(y * components.y + ny);
                    const auto colorValue = values.at(nx + ny * components.x);

                    linearSumR += colorValue.r * basis;
                    linearSumG += colorValue.g * basis;
                    linearSumB += colorValue.b * basis;
                }
            }

            auto linearColor = QColor::fromRgbF(linearSumR, linearSumG, linearSumB);
            image.setPixelColor(x, y, linearColor);
        }
    }

    return image.colorTransformed(fromLinearSRGB);
}

std::optional<int> BlurHash::decode83(const QString &encodedString)
{
    int temp = 0;
    for (const QChar c : std::as_const(encodedString)) {
        const auto index = b83Characters.indexOf(c);
        if (index == -1) {
            return std::nullopt;
        }

        temp = temp * 83 + static_cast<int>(index);
    }

    return temp;
}

BlurHash::Components BlurHash::unpackComponents(const int packedComponents)
{
    return {packedComponents % 9 + 1, packedComponents / 9 + 1};
}

float BlurHash::decodeMaxAC(const int value)
{
    return static_cast<float>(value + 1) / 166.f;
}

QColor BlurHash::decodeAverageColor(const int encodedValue)
{
    const int intR = encodedValue >> 16;
    const int intG = (encodedValue >> 8) & 255;
    const int intB = encodedValue & 255;

    return QColor::fromRgb(intR, intG, intB);
}

float BlurHash::signPow(const float value, const float exp)
{
    return std::copysign(std::pow(std::abs(value), exp), value);
}

BlurHash::ColorF BlurHash::decodeAC(const int value, const float maxAC)
{
    const auto quantR = value / (19 * 19);
    const auto quantG = (value / 19) % 19;
    const auto quantB = value % 19;

    return {.r = signPow((static_cast<float>(quantR) - 9) / 9, 2) * maxAC,
            .g = signPow((static_cast<float>(quantG) - 9) / 9, 2) * maxAC,
            .b = signPow((static_cast<float>(quantB) - 9) / 9, 2) * maxAC};
}

QList<float> BlurHash::calculateWeights(const qsizetype dimension, const qsizetype components)
{
    QList<float> bases(dimension * components, 0.0f);

    const auto scale = static_cast<float>(M_PI) / static_cast<float>(dimension);
    for (qsizetype x = 0; x < dimension; x++) {
        for (qsizetype nx = 0; nx < components; nx++) {
            bases[x * components + nx] = std::cos(scale * static_cast<float>(nx * x));
        }
    }
    return bases;
}
