// Copyright 2016 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "src/observation_store/envelope_maker.h"

#include <utility>

#include <gtest/gtest.h>

#include "src/logging.h"
#include "src/observation_store/observation_store.h"
#include "src/public/lib/status_codes.h"

namespace cobalt::observation_store {

namespace {

constexpr uint32_t kCustomerId = 1;
constexpr uint32_t kProjectId = 1;
constexpr uint32_t kFirstMetricId = 1;
constexpr uint32_t kSecondMetricId = 2;
constexpr uint32_t kThirdMetricId = 3;

// This is the day index for Friday Dec 2, 2016
const uint32_t kUtcDayIndex = 17137;
}  // namespace

class EnvelopeMakerTest : public ::testing::Test {
 public:
  EnvelopeMakerTest()
      : envelope_maker_(new EnvelopeMaker()),
        encrypt_(util::EncryptedMessageMaker::MakeUnencrypted()) {}

  // Returns the current value of envelope_maker_ and resets envelope_maker_
  // to a new EnvelopeMaker constructed using the given optional arguments.
  std::unique_ptr<EnvelopeMaker> ResetEnvelopeMaker(size_t max_bytes_each_observation = SIZE_MAX,
                                                    size_t max_num_bytes = SIZE_MAX) {
    std::unique_ptr<EnvelopeMaker> return_val = std::move(envelope_maker_);
    envelope_maker_ = std::make_unique<EnvelopeMaker>(max_bytes_each_observation, max_num_bytes);
    return return_val;
  }

  // Adds an Observation
  // (as represented by an EncryptedMessage and an ObservationMetadata) of size
  // |num_bytes|, with the specified |metric_id|, to the EnvelopeMaker.
  //
  // |expected_num_batches|is the expected number of batches in the
  // EnvelopeMaker after the new Observation has been added.
  //
  // |expected_this_batch_index| The expected index of the batch that the
  // new Observation should have been put into.
  //
  // |expected_this_batch_size| The expected number of Observations in the
  // batch, including the new Observation.
  void AddObservation(size_t num_bytes, uint32_t metric_id, int expected_num_batches,
                      int expected_this_batch_index, int expected_this_batch_size,
                      StatusCode expected_status_code,
                      const std::string& contribution_id = std::string()) {
    // Create an EncryptedMessage and an ObservationMetadata
    auto message = std::make_unique<EncryptedMessage>();
    // We subtract 1 from |num_bytes| because EnvelopeMaker adds one
    // to its definition of size.
    CHECK(num_bytes > 4);
    message->set_ciphertext(std::string(num_bytes - 4, 'x'));
    auto metadata = std::make_unique<ObservationMetadata>();
    metadata->set_customer_id(kCustomerId);
    metadata->set_project_id(kProjectId);
    metadata->set_metric_id(metric_id);
    metadata->set_day_index(kUtcDayIndex);

    ASSERT_NE(nullptr, envelope_maker_);
    // Add the Observation to the EnvelopeMaker
    size_t size_before_add = envelope_maker_->Size();
    auto m = std::make_unique<StoredObservation>();
    m->set_contribution_id(contribution_id);
    m->set_allocated_encrypted(message.release());
    ASSERT_EQ(expected_status_code,
              envelope_maker_->StoreObservation(std::move(m), std::move(metadata)).error_code());
    size_t size_after_add = envelope_maker_->Size();
    size_t expected_size_change = (expected_status_code == StatusCode::OK ? num_bytes : 0);
    EXPECT_EQ(expected_size_change, size_after_add - size_before_add);

    // Check the number of batches currently in the envelope.
    ASSERT_EQ(expected_num_batches, envelope_maker_->GetEnvelope(encrypt_.get()).batch_size());

    if (expected_status_code != StatusCode::OK) {
      return;
    }

    // Check the ObservationMetadata of the expected batch.
    const auto& batch =
        envelope_maker_->GetEnvelope(encrypt_.get()).batch(expected_this_batch_index);
    EXPECT_EQ(kCustomerId, batch.meta_data().customer_id());
    EXPECT_EQ(kProjectId, batch.meta_data().project_id());
    EXPECT_EQ(metric_id, batch.meta_data().metric_id());
    EXPECT_EQ(kUtcDayIndex, batch.meta_data().day_index());

    // Check the size of the expected batch.
    ASSERT_EQ(expected_this_batch_size, batch.encrypted_observation_size())
        << "batch_index=" << expected_this_batch_index << "; metric_id=" << metric_id;

    // The Observation we just added should be the last one in the batch.
    EXPECT_EQ(num_bytes - 4,
              batch.encrypted_observation(expected_this_batch_size - 1).ciphertext().size());
  }

  // Adds multiple observations to the EnvelopeMaker for the given
  // metric_id. The Observations will have sizes 10 + i for i in [first, limit).
  //
  // expected_num_batches: How many batches do we expecte the EnvelopeMaker to
  // contain after the first add.
  //
  // expected_this_batch_index: Which batch index do we expect this add to
  // have gone into.
  //
  // expected_this_batch_size: What is the expected number of Observations in
  //  the current batch *before* the first add.
  void AddManyObservations(int first, int limit, uint32_t metric_id, int expected_num_batches,
                           int expected_this_batch_index, int expected_this_batch_size) {
    for (int i = first; i < limit; i++) {
      AddObservation(10 + i, metric_id, expected_num_batches, expected_this_batch_index,
                     ++expected_this_batch_size, StatusCode::OK);
    }
  }

  // Adds multiple encoded Observations to two different metrics. Test that
  // the EnvelopeMaker behaves correctly.
  void DoTest() {
    // Add 10 observations for metric 1
    int expected_num_batches = 1;
    int expected_this_batch_index = 0;
    int expected_batch0_size = 0;
    AddManyObservations(10, 20, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                        expected_batch0_size);
    expected_batch0_size += 10;
    // Add 5 more.
    AddManyObservations(15, 20, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                        expected_batch0_size);
    expected_batch0_size += 5;

    // Add 5 observations for metric 2
    expected_num_batches = 2;
    expected_this_batch_index = 1;
    int expected_batch1_size = 0;
    AddManyObservations(15, 20, kSecondMetricId, expected_num_batches, expected_this_batch_index,
                        expected_batch1_size);
    expected_batch1_size += 5;
    // Add 5 more observations for metric 2
    AddManyObservations(20, 25, kSecondMetricId, expected_num_batches, expected_this_batch_index,
                        expected_batch1_size);
    expected_batch1_size += 5;

    // Add 12 more observations for metric 1
    expected_this_batch_index = 0;
    AddManyObservations(11, 23, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                        expected_batch0_size);
    expected_batch0_size += 12;
    // Add 13 more observations for metric 1
    AddManyObservations(11, 24, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                        expected_batch0_size);
    expected_batch0_size += 13;

    // Add 10 more observations for metric 2
    expected_this_batch_index = 1;
    AddManyObservations(10, 20, kSecondMetricId, expected_num_batches, expected_this_batch_index,
                        expected_batch1_size);
    expected_batch1_size += 10;
    // Add 10 more.
    AddManyObservations(10, 20, kSecondMetricId, expected_num_batches, expected_this_batch_index,
                        expected_batch1_size);
    expected_batch1_size += 10;

    const Envelope& envelope = envelope_maker_->GetEnvelope(encrypt_.get());
    EXPECT_EQ(2, envelope.batch_size());
    for (size_t i = 0; i < 2; i++) {
      EXPECT_EQ(i + 1, envelope.batch(i).meta_data().metric_id());
      auto expected_batch_size = (i == 0 ? expected_batch0_size : expected_batch1_size);
      EXPECT_EQ(expected_batch_size, envelope.batch(i).encrypted_observation_size());
    }
  }

  // Returns the ObservationBatch containing the given |metadata|. If
  // this is the first time we have seen the given |metadata| then a
  // new ObservationBatch is created.
  StoredObservationBatch* GetBatch(std::unique_ptr<ObservationMetadata> metadata) {
    return envelope_maker_->GetBatch(std::move(metadata));
  }

 protected:
  std::unique_ptr<EnvelopeMaker> envelope_maker_;
  std::unique_ptr<util::EncryptedMessageMaker> encrypt_;
};

// We perform DoTest() three times with a Clear() between each turn.
// This last tests that Clear() works correctly.
TEST_F(EnvelopeMakerTest, TestThrice) {
  for (int i = 0; i < 3; i++) {
    DoTest();
    envelope_maker_->Clear();
  }
}

// Tests the MergeWith() method.
TEST_F(EnvelopeMakerTest, MergeWith) {
  // Add metric 1 batch to EnvelopeMaker 1 with strings 0..9
  uint32_t metric_id = kFirstMetricId;
  int expected_num_batches = 1;
  int expected_this_batch_index = 0;
  int expected_this_batch_size = 0;
  AddManyObservations(10, 20, metric_id, expected_num_batches, expected_this_batch_index,
                      expected_this_batch_size);

  // Add metric 2 batch to EnvelopeMaker 1 with strings 0..9
  metric_id = kSecondMetricId;
  expected_num_batches = 2;
  expected_this_batch_index = 1;
  AddManyObservations(10, 20, metric_id, expected_num_batches, expected_this_batch_index,
                      expected_this_batch_size);

  // Take EnvelopeMaker 1 and create EnvelopeMaker 2.
  auto envelope_maker1 = ResetEnvelopeMaker();

  // Add metric 2 batch to EnvelopeMaker 2 with strings 10..19
  metric_id = kSecondMetricId;
  expected_num_batches = 1;
  expected_this_batch_index = 0;
  AddManyObservations(10, 20, metric_id, expected_num_batches, expected_this_batch_index,
                      expected_this_batch_size);

  // Add metric 3 to EnvelopeMaker 2 with strings 0..9
  metric_id = kThirdMetricId;
  expected_num_batches = 2;
  expected_this_batch_index = 1;
  AddManyObservations(10, 20, metric_id, expected_num_batches, expected_this_batch_index,
                      expected_this_batch_size);
  // Take EnvelopeMaker 2,
  auto envelope_maker2 = ResetEnvelopeMaker();

  // Now invoke MergeWith to merge EnvelopeMaker 2 into EnvelopeMaker 1.
  envelope_maker1->MergeWith(std::move(envelope_maker2));

  // EnvelopeMaker 2 should be null.
  EXPECT_EQ(envelope_maker2, nullptr);

  // EnvelopeMaker 1 should have three batches for Metrics 1, 2, 3
  EXPECT_FALSE(envelope_maker1->Empty());
  ASSERT_EQ(3, envelope_maker1->GetEnvelope(encrypt_.get()).batch_size());

  // Iterate through each of the batches and check it.
  for (int index = 0; index < 3; index++) {
    // Batch 0 and 2 should have 10 encrypted observations and batch
    // 1 should have 20 because batch 1 from EnvelopeMaker 2 was merged
    // into batch 1 of EnvelopeMaker 1.
    const ObservationBatch& batch = envelope_maker1->GetEnvelope(encrypt_.get()).batch(index);
    EXPECT_EQ(index + 1, batch.meta_data().metric_id());
    auto expected_num_observations = (index == 1 ? 20 : 10);
    ASSERT_EQ(expected_num_observations, batch.encrypted_observation_size());

    // Check each one of the observations.
    for (int i = 0; i < expected_num_observations; i++) {
      // TODO(fxbug.dev/87119): Extract the serialized observation.
    }
  }

  // Now we want to test that after the MergeWith() operation the EnvelopeMaker
  // is still usable. Put EnvelopeMaker 1 back as the test EnvelopeMaker.
  envelope_maker_ = std::move(envelope_maker1);

  // Add string observations 10..19 to metric ID 1 batches 1, 2 and 3.
  for (int metric_id = 1; metric_id <= 3; metric_id++) {
    expected_num_batches = 3;
    expected_this_batch_index = metric_id - 1;
    expected_this_batch_size = (metric_id == 2 ? 20 : 10);
    AddManyObservations(10, 20, metric_id, expected_num_batches, expected_this_batch_index,
                        expected_this_batch_size);
  }
}

// Tests that EnvelopeMaker returns kObservationTooBig when it is supposed to.
TEST_F(EnvelopeMakerTest, ObservationTooBig) {
  // Set max_bytes_each_observation = 105.
  ResetEnvelopeMaker(105);

  int expected_num_batches = 1;
  int expected_this_batch_index = 0;
  int expected_this_batch_size = 1;

  // Add an observation that is not too big.
  AddObservation(105, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                 expected_this_batch_size, StatusCode::OK);

  // Try to add an observation that is too big.
  // We expect the Observation to not be added to the Envelope and so for
  // the Envelope size to not change.
  AddObservation(106, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                 expected_this_batch_size, StatusCode::FAILED_PRECONDITION);

  // Add another observation that is not too big.
  expected_this_batch_size = 2;
  AddObservation(104, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                 expected_this_batch_size, StatusCode::OK);

  // Try to add another observation that is too big.
  // We expect the Observation to not be added to the Envelope and so for
  // the Envelope size to not change.
  AddObservation(107, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                 expected_this_batch_size, StatusCode::FAILED_PRECONDITION);
}

// Tests that EnvelopeMaker returns kStoreFull when it is supposed to.
TEST_F(EnvelopeMakerTest, EnvelopeFull) {
  // Set max_bytes_each_observation = 100, max_num_bytes=1000.
  ResetEnvelopeMaker(100, 1000);

  int expected_this_batch_size = 1;
  int expected_num_batches = 1;
  int expected_this_batch_index = 0;
  for (int i = 0; i < 19; i++) {
    // Add 19 observations of size 50 for a total of 950 bytes
    AddObservation(50, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                   expected_this_batch_size++, StatusCode::OK);
  }
  EXPECT_EQ(950u, envelope_maker_->Size());

  // If we try to add an observation of more than 100 bytes we should
  // get kObservationTooBig.

  // We expect the Observation to not be added to the Envelope and so for
  // the Envelope size to not change.

  AddObservation(101, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                 expected_this_batch_size, StatusCode::FAILED_PRECONDITION);

  // If we try to add an observation of 65 bytes we should
  // get kStoreFull

  AddObservation(65, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                 expected_this_batch_size, StatusCode::RESOURCE_EXHAUSTED);

  // We can add one more observation of size 50.
  AddObservation(50, kFirstMetricId, expected_num_batches, expected_this_batch_index,
                 expected_this_batch_size, StatusCode::OK);
}

TEST_F(EnvelopeMakerTest, CanWriteUnencrypted) {
  auto observation = std::make_unique<StoredObservation>();
  Observation* unencrypted = observation->mutable_unencrypted();
  unencrypted->set_random_id("test123");
  unencrypted->mutable_basic_rappor()->set_data("test");
  observation->set_contribution_id("12345678");

  auto metadata = std::make_unique<ObservationMetadata>();
  metadata->set_customer_id(kCustomerId);
  metadata->set_project_id(kProjectId);
  metadata->set_metric_id(10);

  ASSERT_EQ(
      StatusCode::OK,
      envelope_maker_->StoreObservation(std::move(observation), std::move(metadata)).error_code());
}

TEST_F(EnvelopeMakerTest, CanReadUnencrypted) {
  auto observation = std::make_unique<Observation>();
  observation->set_random_id("test123");
  observation->mutable_basic_rappor()->set_data("test");

  auto encrypted_obs = std::make_unique<EncryptedMessage>();
  encrypt_->Encrypt(*observation, encrypted_obs.get());

  // Verify that our encrypted observation is non-trivial.
  ASSERT_GT(encrypted_obs->ciphertext().size(), 0);

  auto metadata = std::make_unique<ObservationMetadata>();
  metadata->set_customer_id(kCustomerId);
  metadata->set_project_id(kProjectId);
  metadata->set_metric_id(10);

  GetBatch(std::move(metadata))->add_observation()->mutable_unencrypted()->Swap(observation.get());

  auto read_env = envelope_maker_->GetEnvelope(encrypt_.get());
  ASSERT_EQ(read_env.batch_size(), 1);
  ASSERT_EQ(read_env.batch(0).encrypted_observation_size(), 1);

  // Verify that we got the observation we expected.
  ASSERT_EQ(read_env.batch(0).encrypted_observation(0).ciphertext(), encrypted_obs->ciphertext());
}

TEST_F(EnvelopeMakerTest, CanReadWriteUnencrypted) {
  auto observation = std::make_unique<StoredObservation>();
  Observation* unencrypted = observation->mutable_unencrypted();
  observation->set_contribution_id("123456789");
  unencrypted->set_random_id("test123");
  unencrypted->mutable_basic_rappor()->set_data("test");

  auto encrypted_obs = std::make_unique<EncryptedMessage>();
  encrypt_->Encrypt(*unencrypted, encrypted_obs.get());

  // Verify that our encrypted observation is non-trivial.
  ASSERT_GT(encrypted_obs->ciphertext().size(), 0);

  auto metadata = std::make_unique<ObservationMetadata>();
  metadata->set_customer_id(kCustomerId);
  metadata->set_project_id(kProjectId);
  metadata->set_metric_id(10);

  ASSERT_EQ(
      StatusCode::OK,
      envelope_maker_->StoreObservation(std::move(observation), std::move(metadata)).error_code());

  auto read_env = envelope_maker_->GetEnvelope(encrypt_.get());
  ASSERT_EQ(read_env.batch_size(), 1);
  ASSERT_EQ(read_env.batch(0).encrypted_observation_size(), 1);

  // Verify that we got the observation we expected.
  ASSERT_EQ(read_env.batch(0).encrypted_observation(0).ciphertext(), encrypted_obs->ciphertext());
  ASSERT_EQ(read_env.batch(0).encrypted_observation(0).contribution_id(), "123456789");
}

}  // namespace cobalt::observation_store