// Copyright 2020 Google LLC. All Rights Reserved.

#include <stddef.h>
#include <stdint.h>

#include <vector>

#include "api/aead_whitebox.h"
#include "api/test_data.h"
#include "base/logging.h"
#include "benchmarking/data_source.h"
#include "benchmarking/measurements.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace widevine {
namespace {
constexpr size_t kSafeOverhead = 128;
}  // namespace

// Param:
//  - The number of bytes that should be pushed through Decrypt() each time.
//  - The number of iterations to run.
class AeadWhiteboxBenchmark
    : public ::testing::Test,
      public testing::WithParamInterface<std::tuple<size_t, size_t>> {
 protected:
  void SetUp() override {
    const std::vector<uint8_t> context = data_source_.Get(16);
    const std::vector<uint8_t> init_data = GetValidAeadInitData();
    ASSERT_EQ(WB_RESULT_OK,
              WB_Aead_Create(init_data.data(), init_data.size(), context.data(),
                             context.size(), &whitebox_));

    std::tie(plaintext_size_, iterations_) = GetParam();
    ciphertext_size_ = plaintext_size_ + kSafeOverhead;
  }

  void TearDown() override { WB_Aead_Delete(whitebox_); }

  void InitWithCiphertextAsInput() {
    // Make sure both buffers are the correct size for the test now that we know
    // that the ciphertext is the input and the plaintext is the output.
    input_.resize(ciphertext_size_);
    output_.resize(plaintext_size_);

    const auto plaintext = data_source_.Get(plaintext_size_);

    size_t size = ciphertext_size_;
    ASSERT_EQ(WB_RESULT_OK,
              WB_Aead_Encrypt(whitebox_, plaintext.data(), plaintext.size(),
                              input_.data(), &size));
    input_.resize(size);
  }

  void InitWithPlaintextAsInput() {
    // Make sure both buffers are the correct size for the test now that we know
    // that the plaintext is the input and the ciphertext is the output.
    input_ = data_source_.Get(plaintext_size_);
    output_.resize(ciphertext_size_);
  }

  WB_Aead_Whitebox* whitebox_ = nullptr;

  DataSource data_source_;

  size_t plaintext_size_;
  size_t ciphertext_size_;

  std::vector<uint8_t> input_;
  std::vector<uint8_t> output_;

  size_t iterations_;
};

TEST_P(AeadWhiteboxBenchmark, EncryptDuration) {
  InitWithPlaintextAsInput();

  Timer timer;
  Sampler sampler;

  for (size_t i = 0; i < iterations_; i++) {
    timer.Reset();

    size_t output_size = output_.size();
    ASSERT_EQ(WB_RESULT_OK,
              WB_Aead_Encrypt(whitebox_, input_.data(), input_.size(),
                              output_.data(), &output_size));

    sampler.Push(timer.Get());
  }

  PrettyPrint("AEAD Encrypt Duration", sampler, input_.size());
}

TEST_P(AeadWhiteboxBenchmark, EncryptThroughput) {
  InitWithPlaintextAsInput();

  Timer timer;
  timer.Reset();

  for (size_t i = 0; i < iterations_; i++) {
    size_t output_size = output_.size();
    ASSERT_EQ(WB_RESULT_OK,
              WB_Aead_Encrypt(whitebox_, input_.data(), input_.size(),
                              output_.data(), &output_size));
  }

  Throughput throughput(timer.Get(), input_.size() * iterations_);
  PrettyPrint("AEAD Encrypt Throughput", throughput, input_.size());
}

TEST_P(AeadWhiteboxBenchmark, DecryptDuration) {
  InitWithCiphertextAsInput();

  Timer timer;
  Sampler sampler;

  for (size_t i = 0; i < iterations_; i++) {
    timer.Reset();

    size_t output_size = output_.size();
    ASSERT_EQ(WB_RESULT_OK,
              WB_Aead_Decrypt(whitebox_, input_.data(), input_.size(),
                              output_.data(), &output_size));

    sampler.Push(timer.Get());
  }

  PrettyPrint("AEAD Decrypt Duration", sampler, input_.size());
}

TEST_P(AeadWhiteboxBenchmark, DecryptThroughput) {
  InitWithCiphertextAsInput();

  Timer timer;
  timer.Reset();

  for (size_t i = 0; i < iterations_; i++) {
    size_t output_size = output_.size();
    ASSERT_EQ(WB_RESULT_OK,
              WB_Aead_Decrypt(whitebox_, input_.data(), input_.size(),
                              output_.data(), &output_size));
  }

  Throughput throughput(timer.Get(), iterations_ * input_.size());
  PrettyPrint("AEAD Encrypt Throughput", throughput, input_.size());
}

// Run the test will different sizes of input-per-call. The values we have
// picked are divided into three groups, "bytes" and "kilobytes". The services
// team says that a license will normally by 1KB as a base size and then 50 to
// 100 bytes per key.
INSTANTIATE_TEST_SUITE_P(DifferentBytesPerCall,
                         AeadWhiteboxBenchmark,
                         ::testing::Values(std::make_tuple(1024, 100),
                                           std::make_tuple(16 * 1024, 50),
                                           std::make_tuple(256 * 1024, 25)));
}  // namespace widevine