// Copyright 2017 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "components/safe_browsing/core/common/visual_utils.h"

#include <array>

#include "base/containers/span.h"
#include "base/test/scoped_feature_list.h"
#include "base/test/test_discardable_memory_allocator.h"
#include "components/safe_browsing/core/common/features.h"
#include "testing/gmock/include/gmock/gmock.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/skia/include/core/SkColorSpace.h"
#include "third_party/skia/include/private/chromium/SkPMColor.h"
#include "ui/gfx/geometry/rect.h"

namespace safe_browsing::visual_utils {

namespace {

const SkPMColor kSkPMRed = SkPMColorSetARGB(255, 255, 0, 0);
const SkPMColor kSkPMGreen = SkPMColorSetARGB(255, 0, 255, 0);
const SkPMColor kSkPMBlue = SkPMColorSetARGB(255, 0, 0, 255);

}  // namespace

using ::testing::FloatEq;

class VisualUtilsTest : public testing::Test {
 protected:
  void SetUp() override {
    base::DiscardableMemoryAllocator::SetInstance(&test_allocator_);

    sk_sp<SkColorSpace> rec2020 = SkColorSpace::MakeRGB(
        {2.22222f, 0.909672f, 0.0903276f, 0.222222f, 0.0812429f, 0, 0},
        SkNamedGamut::kRec2020);
    SkImageInfo bitmap_info = SkImageInfo::MakeN32(
        1000, 1000, SkAlphaType::kUnpremul_SkAlphaType, rec2020);

    ASSERT_TRUE(bitmap_.tryAllocPixels(bitmap_info));
  }

  void TearDown() override {
    base::DiscardableMemoryAllocator::SetInstance(nullptr);
  }

  void ExpectBlurredImageOneColor(SkColor expected_color,
                                  const VisualFeatures::BlurredImage& image) {
    ExpectBlurredImagePixels(expected_color, image,
                             gfx::Rect(image.width(), image.height()));
  }

  void ExpectBlurredImagePixels(SkColor expected_color,
                                const VisualFeatures::BlurredImage& image,
                                gfx::Rect area) {
    ExpectPixels(expected_color, image.data(), area, image.width());
  }

  void ExpectPixels(SkColor expected_color,
                    const std::string& data_to_check,
                    gfx::Rect area,
                    int width) {
    ASSERT_LE(area.right(), width);
    size_t bottom_right_index =
        ((area.bottom() - 1) * width + area.right() - 1) * 3;
    ASSERT_LT(bottom_right_index, data_to_check.size());
    for (int x = area.x(); x < area.right(); ++x) {
      for (int y = area.y(); y < area.bottom(); ++y) {
        size_t index = (y * width + x) * 3;
        ExpectPixel(expected_color, data_to_check, x, y, index);
      }
    }
  }

  void ExpectPixel(SkColor expected_color,
                   const std::string& data_to_check,
                   size_t x,
                   size_t y,
                   size_t index) {
    ASSERT_LT(index + 2, data_to_check.size());
    EXPECT_EQ(SkColorGetR(expected_color),
              static_cast<unsigned char>(data_to_check[index]))
        << "R component of pixel at x " << x << " and y " << y
        << " is incorrect.";
    EXPECT_EQ(SkColorGetG(expected_color),
              static_cast<unsigned char>(data_to_check[index + 1]))
        << "G component of pixel at x " << x << " and y " << y
        << " is incorrect.";
    EXPECT_EQ(SkColorGetB(expected_color),
              static_cast<unsigned char>(data_to_check[index + 2]))
        << "B component of pixel at x " << x << " and y " << y
        << " is incorrect.";
  }

  // A test bitmap to work with. Initialized to be 1000x1000 in the Rec 2020
  // color space.
  SkBitmap bitmap_;

 private:
  // A DiscardableMemoryAllocator is needed for certain Skia operations.
  base::TestDiscardableMemoryAllocator test_allocator_;
};

TEST_F(VisualUtilsTest, BlurImageWhite) {
  VisualFeatures::BlurredImage blurred;

  // Draw white over the image
  bitmap_.erase(SK_ColorWHITE, SkIRect::MakeXYWH(0, 0, 1000, 1000));

  ASSERT_TRUE(GetBlurredImage(bitmap_, &blurred));

#if BUILDFLAG(IS_ANDROID) || BUILDFLAG(IS_IOS)
  const int kExpectedBlurredWidth = 18;
  const int kExpectedBlurredHeight = 32;
#else
  const int kExpectedBlurredWidth = 48;
  const int kExpectedBlurredHeight = 48;
#endif
  ASSERT_EQ(kExpectedBlurredWidth, blurred.width());
  ASSERT_EQ(kExpectedBlurredHeight, blurred.height());
  ASSERT_EQ(3u * kExpectedBlurredWidth * kExpectedBlurredHeight,
            blurred.data().size());
  ExpectBlurredImageOneColor(SK_ColorWHITE, blurred);
}

TEST_F(VisualUtilsTest, BlurImageRed) {
  VisualFeatures::BlurredImage blurred;

  // Draw red over the image.
  for (int x = 0; x < 1000; x++) {
    for (int y = 0; y < 1000; y++) {
      *bitmap_.getAddr32(x, y) = kSkPMRed;
    }
  }

  ASSERT_TRUE(GetBlurredImage(bitmap_, &blurred));
#if BUILDFLAG(IS_ANDROID) || BUILDFLAG(IS_IOS)
  const int kExpectedBlurredWidth = 18;
  const int kExpectedBlurredHeight = 32;
#else
  const int kExpectedBlurredWidth = 48;
  const int kExpectedBlurredHeight = 48;
#endif
  ASSERT_EQ(kExpectedBlurredWidth, blurred.width());
  ASSERT_EQ(kExpectedBlurredHeight, blurred.height());
  ASSERT_EQ(3u * kExpectedBlurredWidth * kExpectedBlurredHeight,
            blurred.data().size());
  ExpectBlurredImageOneColor(SK_ColorRED, blurred);
}

TEST_F(VisualUtilsTest, BlurImageHalfWhiteHalfBlack) {
  VisualFeatures::BlurredImage blurred;

  // Draw black over half the image.
  bitmap_.erase(SK_ColorBLACK, SkIRect::MakeXYWH(0, 0, 1000, 500));

  // Draw white over half the image
  bitmap_.erase(SK_ColorWHITE, SkIRect::MakeXYWH(0, 500, 1000, 1000));

  ASSERT_TRUE(GetBlurredImage(bitmap_, &blurred));
#if BUILDFLAG(IS_ANDROID) || BUILDFLAG(IS_IOS)
  ASSERT_EQ(18, blurred.width());
  ASSERT_EQ(32, blurred.height());
  ASSERT_EQ(3u * 18u * 32u, blurred.data().size());
  // The middle blocks may have been blurred to something between white and
  // black, so only verify the first 14 and last 14 rows.
  ExpectBlurredImagePixels(SK_ColorBLACK, blurred, gfx::Rect(18, 14));
  ExpectBlurredImagePixels(SK_ColorWHITE, blurred, gfx::Rect(0, 18, 18, 14));
#else
  ASSERT_EQ(48, blurred.width());
  ASSERT_EQ(48, blurred.height());
  ASSERT_EQ(3u * 48u * 48u, blurred.data().size());
  // The middle blocks may have been blurred to something between white and
  // black, so only verify the first 22 and last 22 rows.
  ExpectBlurredImagePixels(SK_ColorBLACK, blurred, gfx::Rect(48, 22));
  ExpectBlurredImagePixels(SK_ColorWHITE, blurred, gfx::Rect(0, 26, 48, 22));
#endif
}

TEST_F(VisualUtilsTest, BlockMeanAverageOneBlock) {
  // Draw black over half the image.
  bitmap_.erase(SK_ColorBLACK, SkIRect::MakeXYWH(0, 0, 1000, 500));

  // Draw white over half the image
  bitmap_.erase(SK_ColorWHITE, SkIRect::MakeXYWH(0, 500, 1000, 1000));

  std::unique_ptr<SkBitmap> blocks = BlockMeanAverage(bitmap_, 1000);
  ASSERT_EQ(1, blocks->width());
  ASSERT_EQ(1, blocks->height());
  EXPECT_EQ(blocks->getColor(0, 0), SkColorSetRGB(127, 127, 127));
}

TEST_F(VisualUtilsTest, BlockMeanAveragePartialBlocks) {
  // Draw a white, red, green, and blue box with the expected block sizes.
  bitmap_.erase(SK_ColorWHITE, SkIRect::MakeXYWH(0, 0, 600, 600));

  for (int x = 600; x < 1000; x++) {
    for (int y = 0; y < 600; y++) {
      *bitmap_.getAddr32(x, y) = kSkPMRed;
    }
  }

  for (int x = 0; x < 600; x++) {
    for (int y = 600; y < 1000; y++) {
      *bitmap_.getAddr32(x, y) = kSkPMGreen;
    }
  }

  for (int x = 600; x < 1000; x++) {
    for (int y = 600; y < 1000; y++) {
      *bitmap_.getAddr32(x, y) = kSkPMBlue;
    }
  }

  std::unique_ptr<SkBitmap> blocks = BlockMeanAverage(bitmap_, 600);
  ASSERT_EQ(2, blocks->width());
  ASSERT_EQ(2, blocks->height());
  EXPECT_EQ(blocks->getColor(0, 0), SK_ColorWHITE);

  EXPECT_EQ(*blocks->getAddr32(1, 0), kSkPMRed);
  EXPECT_EQ(*blocks->getAddr32(0, 1), kSkPMGreen);
  EXPECT_EQ(*blocks->getAddr32(1, 1), kSkPMBlue);
}

TEST_F(VisualUtilsTest, NonSquareBlurredImage) {
  base::test::ScopedFeatureList feature_list;
  feature_list.InitAndEnableFeatureWithParameters(
      kVisualFeaturesSizes, {{"phash_width", "108"}, {"phash_height", "192"}});

  VisualFeatures::BlurredImage blurred;

  // Draw white over the image
  bitmap_.erase(SK_ColorWHITE, SkIRect::MakeXYWH(0, 0, 1000, 1000));

  ASSERT_TRUE(GetBlurredImage(bitmap_, &blurred));
  ASSERT_EQ(18, blurred.width());
  ASSERT_EQ(32, blurred.height());
  ASSERT_EQ(3u * 18u * 32u, blurred.data().size());
  for (size_t i = 0; i < 18u * 32u; i++) {
    EXPECT_EQ('\xff', blurred.data()[3 * i]);
    EXPECT_EQ('\xff', blurred.data()[3 * i + 1]);
    EXPECT_EQ('\xff', blurred.data()[3 * i + 2]);
  }
}

TEST_F(VisualUtilsTest, EncodeScreenshot) {
  const int kBitmapWidth = 40;
  const int kBitmapHeight = 40;
  SkBitmap bitmap;
  bitmap.allocN32Pixels(kBitmapWidth, kBitmapHeight);
  bitmap.eraseColor(SK_ColorWHITE);
  bitmap.erase(SK_ColorRED, SkIRect::MakeXYWH(0, 0, 20, 10));
  bitmap.erase(SK_ColorBLUE, SkIRect::MakeXYWH(20, 0, 20, 10));

  VisualFeatures::Screenshot encoded_screenshot;
  EncodeScreenshot(bitmap, &encoded_screenshot);
  EXPECT_EQ(kBitmapWidth, encoded_screenshot.width());
  EXPECT_EQ(kBitmapHeight, encoded_screenshot.height());
  ExpectPixels(SK_ColorRED, encoded_screenshot.data(), gfx::Rect(20, 10),
               kBitmapWidth);
  ExpectPixels(SK_ColorBLUE, encoded_screenshot.data(),
               gfx::Rect(20, 0, 20, 10), kBitmapWidth);
  ExpectPixels(SK_ColorWHITE, encoded_screenshot.data(),
               gfx::Rect(0, 10, 40, 30), kBitmapWidth);
}

}  // namespace safe_browsing::visual_utils