/*
 * Copyright 2013 Google Inc.
 *
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#include "include/codec/SkCodec.h"
#include "include/core/SkPixmap.h"
#include "include/core/SkStream.h"
#include "include/core/SkYUVASizeInfo.h"
#include "include/private/SkTemplates.h"
#include "src/core/SkAutoMalloc.h"
#include "tests/Test.h"
#include "tools/Resources.h"

static void codec_yuv(skiatest::Reporter* reporter,
                      const char path[],
                      const SkYUVAInfo* expectedInfo) {
    std::unique_ptr<SkStream> stream(GetResourceAsStream(path));
    if (!stream) {
        return;
    }
    std::unique_ptr<SkCodec> codec(SkCodec::MakeFromStream(std::move(stream)));
    REPORTER_ASSERT(reporter, codec);
    if (!codec) {
        return;
    }

    // Test queryYUBAInfo()
    SkYUVAInfo yuvaInfo;
    static_assert(SkYUVAInfo::kMaxPlanes == 4);
    SkColorType colorTypes[SkYUVAInfo::kMaxPlanes] = {kUnknown_SkColorType,
                                                      kUnknown_SkColorType,
                                                      kUnknown_SkColorType,
                                                      kUnknown_SkColorType};
    size_t rowBytes[SkYUVAInfo::kMaxPlanes] = {};

    // All params are required to be non-null.
    bool success = codec->queryYUVAInfo(&yuvaInfo, colorTypes, nullptr);
    REPORTER_ASSERT(reporter, !success);
    success = codec->queryYUVAInfo(&yuvaInfo, nullptr, rowBytes);
    REPORTER_ASSERT(reporter, !success);
    success = codec->queryYUVAInfo(nullptr, colorTypes, rowBytes);
    REPORTER_ASSERT(reporter, !success);
    success = codec->queryYUVAInfo(&yuvaInfo, colorTypes, rowBytes);
    REPORTER_ASSERT(reporter, SkToBool(expectedInfo) == success);
    if (!success) {
        return;
    }
    REPORTER_ASSERT(reporter, *expectedInfo == yuvaInfo);

    SkImageInfo ii[SkYUVAInfo::kMaxPlanes];

    SkISize planeDims[SkYUVAInfo::kMaxPlanes] = {};
    int numPlanes = yuvaInfo.expectedPlaneDims(planeDims);
    REPORTER_ASSERT(reporter, numPlanes <= SkYUVAInfo::kMaxPlanes);
    size_t totalBytes = 0;
    for (int i = 0; i < numPlanes; ++i) {
        REPORTER_ASSERT(reporter, !planeDims[i].isEmpty());
        REPORTER_ASSERT(reporter, colorTypes[i] != kUnknown_SkColorType);
        ii[i] = SkImageInfo::Make(planeDims[i], colorTypes[i], kPremul_SkAlphaType);
        REPORTER_ASSERT(reporter, ii[i].validRowBytes(rowBytes[i]));
        totalBytes += ii[i].height()*rowBytes[i];
    }
    for (int i = numPlanes; i < SkYUVAInfo::kMaxPlanes; ++i) {
        REPORTER_ASSERT(reporter, planeDims[i].isZero());
        REPORTER_ASSERT(reporter, colorTypes[i] == kUnknown_SkColorType);
        REPORTER_ASSERT(reporter, rowBytes[i] == 0);
    }

    // Allocate the memory for the YUV decode.
    SkAutoMalloc storage(totalBytes);
    SkPixmap planes[SkYUVAInfo::kMaxPlanes];
    char* addr = static_cast<char*>(storage.get());
    for (int i = 0; i < numPlanes; ++i) {
        planes[i].reset(ii[i], addr, rowBytes[i]);
        addr += planes[i].height()*rowBytes[i];
    }

    // Test getYUVAPlanes()
    REPORTER_ASSERT(reporter, SkCodec::kInvalidInput == codec->getYUVAPlanes(nullptr));
    REPORTER_ASSERT(reporter, SkCodec::kSuccess == codec->getYUVAPlanes(planes));
}

DEF_TEST(Jpeg_YUV_Codec, r) {
    auto setExpectations = [](SkISize dims, SkYUVAInfo::PlanarConfig planarConfig) {
        return SkYUVAInfo(dims,
                          planarConfig,
                          kJPEG_Full_SkYUVColorSpace,
                          kTopLeft_SkEncodedOrigin,
                          SkYUVAInfo::Siting::kCentered,
                          SkYUVAInfo::Siting::kCentered);
    };

    SkYUVAInfo expectations = setExpectations({128, 128}, SkYUVAInfo::PlanarConfig::kY_U_V_420);
    codec_yuv(r, "images/color_wheel.jpg", &expectations);

    // H2V2
    expectations = setExpectations({512, 512}, SkYUVAInfo::PlanarConfig::kY_U_V_420);
    codec_yuv(r, "images/mandrill_512_q075.jpg", &expectations);

    // H1V1
    expectations = setExpectations({512, 512}, SkYUVAInfo::PlanarConfig::kY_U_V_444);
    codec_yuv(r, "images/mandrill_h1v1.jpg", &expectations);

    // H2V1
    expectations = setExpectations({512, 512}, SkYUVAInfo::PlanarConfig::kY_U_V_422);
    codec_yuv(r, "images/mandrill_h2v1.jpg", &expectations);

    // Non-power of two dimensions
    expectations = setExpectations({439, 154}, SkYUVAInfo::PlanarConfig::kY_U_V_420);
    codec_yuv(r, "images/cropped_mandrill.jpg", &expectations);

    expectations = setExpectations({8, 8}, SkYUVAInfo::PlanarConfig::kY_U_V_420);
    codec_yuv(r, "images/randPixels.jpg", &expectations);

    // Progressive images
    expectations = setExpectations({512, 512}, SkYUVAInfo::PlanarConfig::kY_U_V_444);
    codec_yuv(r, "images/brickwork-texture.jpg", &expectations);
    codec_yuv(r, "images/brickwork_normal-map.jpg", &expectations);

    // A CMYK encoded image should fail.
    codec_yuv(r, "images/CMYK.jpg", nullptr);
    // A grayscale encoded image should fail.
    codec_yuv(r, "images/grayscale.jpg", nullptr);
    // A PNG should fail.
    codec_yuv(r, "images/arrow.png", nullptr);
}

#include "include/effects/SkColorMatrix.h"
#include "src/core/SkYUVMath.h"

// Be sure that the two matrices are inverses of each other
// (i.e. rgb2yuv and yuv2rgb
DEF_TEST(YUVMath, reporter) {
    const SkYUVColorSpace spaces[] = {
        kJPEG_SkYUVColorSpace,
        kRec601_SkYUVColorSpace,
        kRec709_SkYUVColorSpace,
        kBT2020_SkYUVColorSpace,
        kIdentity_SkYUVColorSpace,
    };

    // Not sure what the theoretical precision we can hope for is, so pick a big value that
    // passes (when I think we're correct).
    const float tolerance = 1.0f/(1 << 18);

    for (auto cs : spaces) {
        float r2y[20], y2r[20];
        SkColorMatrix_RGB2YUV(cs, r2y);
        SkColorMatrix_YUV2RGB(cs, y2r);

        SkColorMatrix r2ym, y2rm;
        r2ym.setRowMajor(r2y);
        y2rm.setRowMajor(y2r);
        r2ym.postConcat(y2rm);

        float tmp[20];
        r2ym.getRowMajor(tmp);
        for (int i = 0; i < 20; ++i) {
            float expected = 0;
            if (i % 6 == 0) {   // diagonal
                expected = 1;
            }
            REPORTER_ASSERT(reporter, SkScalarNearlyEqual(tmp[i], expected, tolerance));
        }
    }
}