// Copyright (C) 2023 The Qt Company Ltd.
// SPDX-License-Identifier: LicenseRef-Qt-Commercial OR GPL-3.0-only

#include <QtTest/QtTest>

#include <qvideoframe.h>
#include <qvideoframeformat.h>
#include "private/qmemoryvideobuffer_p.h"
#include "private/qplatformmediaintegration_p.h"
#include "private/qimagevideobuffer_p.h"
#include <QtGui/QColorSpace>
#include <QtGui/QImage>
#include <QtCore/QPointer>

QT_USE_NAMESPACE

namespace {

struct TestParams
{
    QString fileName;
    QVideoFrameFormat::PixelFormat pixelFormat;
    QVideoFrameFormat::ColorSpace colorSpace;
    QVideoFrameFormat::ColorRange colorRange;
};

QString toString(QVideoFrameFormat::ColorRange r)
{
    switch (r) {
    case QVideoFrameFormat::ColorRange_Video:
        return "Video";
    case QVideoFrameFormat::ColorRange_Full:
        return "Full";
    default:
        Q_ASSERT(false);
        return "";
    }
}

std::vector<QVideoFrameFormat::ColorRange> colorRanges()
{
    return {
        QVideoFrameFormat::ColorRange_Video,
        QVideoFrameFormat::ColorRange_Full,
    };
}

// clang-format off

static const QHash<QVideoFrameFormat::PixelFormat, const char*> s_formats {
    { QVideoFrameFormat::Format_ARGB8888, "argb8888" },
    { QVideoFrameFormat::Format_ARGB8888_Premultiplied, "argb8888_premultiplied" },
    { QVideoFrameFormat::Format_XRGB8888, "xrgb8888" },
    { QVideoFrameFormat::Format_BGRA8888, "bgra8888" },
    { QVideoFrameFormat::Format_BGRA8888_Premultiplied, "bgra8888_premultiplied" },
    { QVideoFrameFormat::Format_BGRX8888, "bgrx8888" },
    { QVideoFrameFormat::Format_ABGR8888, "abgr8888" },
    { QVideoFrameFormat::Format_XBGR8888, "xbgr8888" },
    { QVideoFrameFormat::Format_RGBA8888, "rgba8888" },
    { QVideoFrameFormat::Format_RGBX8888, "rgbx8888" },
    { QVideoFrameFormat::Format_NV12, "nv12" },
    { QVideoFrameFormat::Format_NV21, "nv21" },
    { QVideoFrameFormat::Format_IMC1, "imc1" },
    { QVideoFrameFormat::Format_IMC2, "imc2" },
    { QVideoFrameFormat::Format_IMC3, "imc3" },
    { QVideoFrameFormat::Format_IMC4, "imc4" },
    //{ QVideoFrameFormat::Format_AYUV, "ayuv" }, // TODO: Fixme (No corresponding FFmpeg format available)
    //{ QVideoFrameFormat::Format_AYUV_Premultiplied, "ayuv_premultiplied" }, // TODO: Fixme (No corresponding FFmpeg format available)
    { QVideoFrameFormat::Format_YV12, "yv12" },
    { QVideoFrameFormat::Format_YUV420P, "420p" },
    { QVideoFrameFormat::Format_YUV422P, "422p" },
    { QVideoFrameFormat::Format_UYVY, "uyvy" },
    { QVideoFrameFormat::Format_YUYV, "yuyv" },
    { QVideoFrameFormat::Format_Y8, "y8" },
    { QVideoFrameFormat::Format_Y16, "y16" },
    { QVideoFrameFormat::Format_P010, "p010" },
    { QVideoFrameFormat::Format_P016, "p016" },
    { QVideoFrameFormat::Format_YUV420P10, "yuv420p10" }
};

// clang-format on

QString toString(QVideoFrameFormat::PixelFormat f)
{
    if (!s_formats.contains(f)) {
        Q_ASSERT(false);
        return {};
    }

    return s_formats.value(f);
}

QList<QVideoFrameFormat::PixelFormat> pixelFormats()
{
    return s_formats.keys();
}

bool isSupportedPixelFormat(QVideoFrameFormat::PixelFormat pixelFormat)
{
#ifdef Q_OS_ANDROID
    // TODO: QTBUG-125238
    switch (pixelFormat) {
    case QVideoFrameFormat::Format_Y16:
    case QVideoFrameFormat::Format_P010:
    case QVideoFrameFormat::Format_P016:
    case QVideoFrameFormat::Format_YUV420P10:
        return false;
    default:
        return true;
    }
#else
    return true;
#endif
}


QString toString(QVideoFrameFormat::ColorSpace s)
{
    switch (s) {
    case QVideoFrameFormat::ColorSpace_BT601:
        return "BT601";
    case QVideoFrameFormat::ColorSpace_BT709:
        return "BT709";
    case QVideoFrameFormat::ColorSpace_AdobeRgb:
        return "AdobeRgb";
    case QVideoFrameFormat::ColorSpace_BT2020:
        return "BT2020";
    default:
        Q_ASSERT(false);
        return "";
    }
}

std::vector<QVideoFrameFormat::ColorSpace> colorSpaces()
{
    return { QVideoFrameFormat::ColorSpace_BT601, QVideoFrameFormat::ColorSpace_BT709,
             QVideoFrameFormat::ColorSpace_AdobeRgb, QVideoFrameFormat::ColorSpace_BT2020 };
}

QString name(const TestParams &p)
{
    return QStringLiteral("%1_%2_%3_%4")
            .arg(p.fileName)
            .arg(toString(p.pixelFormat))
            .arg(toString(p.colorSpace))
            .arg(toString(p.colorRange));
}

QString path(const QTemporaryDir &dir, const TestParams &param, const QString &suffix = ".png")
{
    return dir.filePath(name(param) + suffix);
}

QVideoFrame createTestFrame(const TestParams &params, const QImage &image)
{
    QVideoFrameFormat format(image.size(), params.pixelFormat);
    format.setColorRange(params.colorRange);
    format.setColorSpace(params.colorSpace);
    format.setColorTransfer(QVideoFrameFormat::ColorTransfer_Unknown);

    auto buffer = std::make_unique<QImageVideoBuffer>(image);
    QVideoFrameFormat imageFormat = {
        image.size(), QVideoFrameFormat::pixelFormatFromImageFormat(image.format())
    };

    QVideoFrame source{ buffer.release(), imageFormat };
    return QPlatformMediaIntegration::instance()->convertVideoFrame(source, format);
}

struct ImageDiffReport
{
    int DiffCountAboveThreshold; // Number of channel differences above threshold
    int MaxDiff;                 // Maximum difference between two images (max across channels)
    int PixelCount;              // Number of pixels in the image
    QImage DiffImage;            // The difference image (absolute per-channel difference)
};

double aboveThresholdDiffRatio(const ImageDiffReport &report)
{
    return static_cast<double>(report.DiffCountAboveThreshold) / report.PixelCount;
}

int maxChannelDiff(QRgb lhs, QRgb rhs)
{
    // clang-format off
    return std::max({ std::abs(qRed(lhs)   - qRed(rhs)),
                      std::abs(qGreen(lhs) - qGreen(rhs)),
                      std::abs(qBlue(lhs)  - qBlue(rhs)) });
    // clang-format on
}

int clampedAbsDiff(int lhs, int rhs)
{
    return std::clamp(std::abs(lhs - rhs), 0, 255);
}

QRgb pixelDiff(QRgb lhs, QRgb rhs)
{
    return qRgb(clampedAbsDiff(qRed(lhs), qRed(rhs)), clampedAbsDiff(qGreen(lhs), qGreen(rhs)),
                clampedAbsDiff(qBlue(lhs), qBlue(rhs)));
}

std::optional<ImageDiffReport> compareImagesRgb32(const QImage &computed, const QImage &baseline,
                                             int channelThreshold)
{
    Q_ASSERT(baseline.format() == QImage::Format_RGB32);

    if (computed.size() != baseline.size())
        return {};

    if (computed.format() != baseline.format())
        return {};

    if (computed.colorSpace() != baseline.colorSpace())
        return {};

    const QSize size = baseline.size();

    ImageDiffReport report{};
    report.PixelCount = size.width() * size.height();
    report.DiffImage = QImage(size, baseline.format());

    // Iterate over all pixels and update report
    for (int l = 0; l < size.height(); l++) {
        const QRgb *colorComputed = reinterpret_cast<const QRgb *>(computed.constScanLine(l));
        const QRgb *colorBaseline = reinterpret_cast<const QRgb *>(baseline.constScanLine(l));
        QRgb *colorDiff = reinterpret_cast<QRgb *>(report.DiffImage.scanLine(l));

        int w = size.width();
        while (w--) {
            *colorDiff = pixelDiff(*colorComputed, *colorBaseline);
            if (*colorComputed != *colorBaseline) {
                const int diff = maxChannelDiff(*colorComputed, *colorBaseline);

                if (diff > report.MaxDiff)
                    report.MaxDiff = diff;

                if (diff > channelThreshold)
                    ++report.DiffCountAboveThreshold;
            }

            ++colorComputed;
            ++colorBaseline;
            ++colorDiff;
        }
    }
    return report;
}

bool copyAllFiles(const QDir &source, const QDir &dest)
{
    if (!source.exists() || !dest.exists())
        return false;

    QDirIterator it(source);
    while (it.hasNext()) {
        QFileInfo file{ it.next() };
        if (file.isFile()) {
            const QString destination = dest.absolutePath() + "/" + file.fileName();
            QFile::copy(file.absoluteFilePath(), destination);
        }
    }

    return true;
}

class ReferenceData
{
public:
    ReferenceData()
    {
        m_testdataDir = QTest::qExtractTestData("testdata");
        if (!m_testdataDir)
            m_testdataDir = QSharedPointer<QTemporaryDir>(new QTemporaryDir);
    }

    ~ReferenceData()
    {
        if (m_testdataDir->autoRemove())
            return;

        QString resultPath = m_testdataDir->path();
        if (qEnvironmentVariableIsSet("COIN_CTEST_RESULTSDIR")) {
            const QDir sourceDir = m_testdataDir->path();
            const QDir resultsDir{ qEnvironmentVariable("COIN_CTEST_RESULTSDIR") };
            if (!copyAllFiles(sourceDir, resultsDir)) {
                qDebug() << "Failed to copy files to COIN_CTEST_RESULTSDIR";
            } else {
                resultPath = resultsDir.path();
            }
        }

        qDebug() << "Images with differences were found. The output images with differences"
                 << "can be found in" << resultPath << ". Review the images and if the"
                 << "differences are expected, please update the testdata with the new"
                 << "output images";
    }

    QImage getReference(const TestParams &param) const
    {
        const QString referenceName = name(param);
        const QString referencePath = m_testdataDir->filePath(referenceName + ".png");
        QImage result;
        if (result.load(referencePath))
            return result;
        return {};
    }

    void saveNewReference(const QImage &reference, const TestParams &params) const
    {
        const QString filename = path(*m_testdataDir, params);
        if (!reference.save(filename)) {
            qDebug() << "Failed to save reference file";
            Q_ASSERT(false);
        }

        m_testdataDir->setAutoRemove(false);
    }

    bool saveComputedImage(const TestParams &params, const QImage &image, const QString& suffix) const
    {
        if (!image.save(path(*m_testdataDir, params, suffix))) {
            qDebug() << "Unexpectedly failed to save actual image to file";
            Q_ASSERT(false);
            return false;
        }
        m_testdataDir->setAutoRemove(false);
        return true;
    }

    QImage getTestdata(const QString &name)
    {
        const QString filePath = m_testdataDir->filePath(name);
        QImage image;
        if (image.load(filePath))
            return image;
        return {};
    }

private:
    QSharedPointer<QTemporaryDir> m_testdataDir;
};

std::optional<ImageDiffReport> compareToReference(const TestParams &params, const QImage &actual,
                                                  const ReferenceData &references,
                                                  int maxChannelThreshold)
{
    const QImage expected = references.getReference(params);
    if (expected.isNull()) {
        // Reference image does not exist. Create one. Adding this to
        // testdata directory is a manual job.
        references.saveNewReference(actual, params);
        qDebug() << "Reference image is missing. Please update testdata directory with the missing "
                    "reference image";
        return {};
    }

    // Convert to RGB32 to simplify image comparison
    const QImage computed = actual.convertToFormat(QImage::Format_RGB32);
    const QImage baseline = expected.convertToFormat(QImage::Format_RGB32);

    std::optional<ImageDiffReport> diffReport = compareImagesRgb32(computed, baseline, maxChannelThreshold);
    if (!diffReport)
        return {};

    if (diffReport->MaxDiff > 0) {
        // Images are not equal, and may require manual inspection
        if (!references.saveComputedImage(params, computed, "_actual.png"))
            return {};
        if (!references.saveComputedImage(params, diffReport->DiffImage, "_diff.png"))
            return {};
    }

    return diffReport;
}

} // namespace

class tst_qvideoframecolormanagement : public QObject
{
    Q_OBJECT
private slots:

    void toImage_savesWithCorrectColors_data()
    {
        QTest::addColumn<QString>("fileName");
        QTest::addColumn<TestParams>("params");
        for (const char *file : { "umbrellas.jpg" }) {
            for (const QVideoFrameFormat::PixelFormat pixelFormat : pixelFormats()) {
                for (const QVideoFrameFormat::ColorSpace colorSpace : colorSpaces()) {
                    for (const QVideoFrameFormat::ColorRange colorRange : colorRanges()) {

                        if (!isSupportedPixelFormat(pixelFormat))
                            continue;

                        TestParams param{ file, pixelFormat, colorSpace, colorRange };
                        QTest::addRow("%s", name(param).toLatin1().data()) << file << param;
                    }
                }
            }
        }
    }

    // This test is a regression test for the QMultimedia display pipeline.
    // It compares rendered output (as created by toImage) against reference
    // images stored to file. The reference images were created by the test
    // itself, and does not verify correctness, just changes to render output.
    void toImage_savesWithCorrectColors()
    {
        QFETCH(const QString, fileName);
        QFETCH(const TestParams, params);

        const QImage templateImage = m_reference.getTestdata(fileName);
        QVERIFY(!templateImage.isNull());

        const QVideoFrame frame = createTestFrame(params, templateImage);

        // Act
        const QImage actual = frame.toImage();

        // Assert
        constexpr int diffThreshold = 4;
        std::optional<ImageDiffReport> result =
                compareToReference(params, actual, m_reference, diffThreshold);

        // Sanity checks
        QVERIFY(result.has_value());
        QCOMPARE_GT(result->PixelCount, 0);

        // Verify that images are similar
        const double ratioAboveThreshold =
                static_cast<double>(result->DiffCountAboveThreshold) / result->PixelCount;

        // These thresholds are empirically determined to allow tests to pass in CI.
        // If tests fail, review the difference between the reference and actual
        // output to determine if it is a platform dependent inaccuracy before
        // adjusting the limits
        QCOMPARE_LT(ratioAboveThreshold, 0.01); // Fraction of pixels with larger differences
        QCOMPARE_LT(result->MaxDiff, 6); // Maximum per-channel difference
    }

private:
    ReferenceData m_reference;
};

QTEST_MAIN(tst_qvideoframecolormanagement)

#include "tst_qvideoframecolormanagement.moc"