#include "src/lib/privacy/private_index_decoding.h"

#include <gmock/gmock.h>
#include <gtest/gtest.h>

#include "src/logger/event_vector_index.h"

namespace cobalt {
namespace {

const int64_t kMinValue = -10;
const int64_t kMaxValue = 50;
const uint64_t kMaxCount = 100;
const uint64_t kNumIndexPoints = 5;
const uint64_t kMaxBucketIndex = 10;

}  // namespace

class PrivateIndexDecodingTest : public testing::Test {
 protected:
  void SetUp() override { PrepareMetricDimensions(); }

  void PrepareMetricDimensions() {
    auto dimension = metric_dimensions_.Add();
    dimension->set_dimension("dimension_consecutive");
    (*dimension->mutable_event_codes())[0] = "dim_0";
    (*dimension->mutable_event_codes())[1] = "dim_1";
    (*dimension->mutable_event_codes())[2] = "dim_2";

    dimension = metric_dimensions_.Add();
    dimension->set_dimension("dimension_non_consecutive");
    (*dimension->mutable_event_codes())[0] = "dim_0";
    (*dimension->mutable_event_codes())[100] = "dim_1";
    (*dimension->mutable_event_codes())[200] = "dim_2";
    (*dimension->mutable_event_codes())[300] = "dim_3";

    dimension = metric_dimensions_.Add();
    dimension->set_dimension("dimension_max_event_code");
    dimension->set_max_event_code(10);
    (*dimension->mutable_event_codes())[1] = "dim_0";
    (*dimension->mutable_event_codes())[2] = "dim_1";
  }

  uint64_t MaxEventVectorIndex() { return logger::GetNumEventVectors(metric_dimensions_) - 1; }

  uint64_t MaxIndexForNumericValue() {
    return logger::GetNumEventVectors(metric_dimensions_) * kNumIndexPoints - 1;
  }

  uint64_t MinIndexForCount() { return MaxIndexForNumericValue() + 1; }

  uint64_t MaxIndexForCount() { return 2 * MinIndexForCount() - 1; }

  uint64_t MaxIndexForHistogramBucket() {
    return logger::GetNumEventVectors(metric_dimensions_) * (kMaxBucketIndex + 1) - 1;
  }

  uint64_t MaxIndexForHistogramBucketAndCount() {
    return logger::GetNumEventVectors(metric_dimensions_) * (kMaxBucketIndex + 1) *
               kNumIndexPoints -
           1;
  }

  google::protobuf::RepeatedPtrField<MetricDefinition::MetricDimension> metric_dimensions_;
};

TEST_F(PrivateIndexDecodingTest, DecodeAsEventVectorValid) {
  std::vector<uint32_t> event_vector;
  // Check that the minimum private index is decoded correctly.
  util::Status status = DecodePrivateIndexAsEventVector(0, metric_dimensions_, &event_vector);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({0, 0, 0}));

  // Check that the maximum private index is decoded correctly.
  status =
      DecodePrivateIndexAsEventVector(MaxEventVectorIndex(), metric_dimensions_, &event_vector);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({2, 300, 10}));
}

TEST_F(PrivateIndexDecodingTest, DecodeAsEventVectorInvalid) {
  std::vector<uint32_t> event_vector;
  util::Status status =
      DecodePrivateIndexAsEventVector(MaxEventVectorIndex() + 1, metric_dimensions_, &event_vector);
  EXPECT_FALSE(status.ok());
}

TEST_F(PrivateIndexDecodingTest, DecodeAsIntegerValid) {
  std::vector<uint32_t> event_vector;
  int64_t integer_value;
  // Check that the minimum private index is decoded correctly.
  util::Status status = DecodePrivateIndexAsInteger(0, metric_dimensions_, kMinValue, kMaxValue,
                                                    kNumIndexPoints, &event_vector, &integer_value);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({0, 0, 0}));
  EXPECT_EQ(integer_value, kMinValue);

  // Check that the maximum private index is decoded correctly.
  status = DecodePrivateIndexAsInteger(MaxIndexForNumericValue(), metric_dimensions_, kMinValue,
                                       kMaxValue, kNumIndexPoints, &event_vector, &integer_value);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({2, 300, 10}));
  EXPECT_EQ(integer_value, kMaxValue);
}

TEST_F(PrivateIndexDecodingTest, DecodeAsIntegerInvalid) {
  std::vector<uint32_t> event_vector;
  int64_t integer_value;
  util::Status status =
      DecodePrivateIndexAsInteger(MaxIndexForNumericValue() + 1, metric_dimensions_, kMinValue,
                                  kMaxValue, kNumIndexPoints, &event_vector, &integer_value);
  EXPECT_FALSE(status.ok());
}

TEST_F(PrivateIndexDecodingTest, DecodeAsDoubleValid) {
  std::vector<uint32_t> event_vector;
  double double_value;
  // Check that the minimum private index is decoded correctly.
  util::Status status = DecodePrivateIndexAsDouble(0, metric_dimensions_, kMinValue, kMaxValue,
                                                   kNumIndexPoints, &event_vector, &double_value);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({0, 0, 0}));
  EXPECT_EQ(double_value, static_cast<double>(kMinValue));

  // Check that an intermediate private index is decoded correctly.
  // The private index (MaxEventVectorIndex() + 1) should be the first index whose |double_value| is
  // larger than |kMinValue|, and its |event_vector| should be the lexicographically smallest event
  // vector.
  status = DecodePrivateIndexAsDouble(MaxEventVectorIndex() + 1, metric_dimensions_, kMinValue,
                                      kMaxValue, kNumIndexPoints, &event_vector, &double_value);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({0, 0, 0}));
  EXPECT_THAT(double_value,
              testing::DoubleEq(kMinValue + static_cast<double>(kMaxValue - kMinValue) /
                                                (kNumIndexPoints - 1)));

  // Check that the maximum private index is decoded correctly.
  status = DecodePrivateIndexAsDouble(MaxIndexForNumericValue(), metric_dimensions_, kMinValue,
                                      kMaxValue, kNumIndexPoints, &event_vector, &double_value);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({2, 300, 10}));
  EXPECT_EQ(double_value, static_cast<double>(kMaxValue));
}

TEST_F(PrivateIndexDecodingTest, DecodeAsDoubleInvalid) {
  std::vector<uint32_t> event_vector;
  double double_value;
  util::Status status =
      DecodePrivateIndexAsDouble(MaxIndexForNumericValue() + 1, metric_dimensions_, kMinValue,
                                 kMaxValue, kNumIndexPoints, &event_vector, &double_value);
  EXPECT_FALSE(status.ok());
}

TEST_F(PrivateIndexDecodingTest, DecodeAsSumOrCountValidSum) {
  std::vector<uint32_t> event_vector;
  SumOrCount sum_or_count;
  // Check that the minimum private index that corresponds to sum is decoded correctly.
  util::Status status =
      DecodePrivateIndexAsSumOrCount(0, metric_dimensions_, kMinValue, kMaxValue, kMaxCount,
                                     kNumIndexPoints, &event_vector, &sum_or_count);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({0, 0, 0}));
  EXPECT_EQ(sum_or_count.type, SumOrCount::SUM);
  EXPECT_EQ(sum_or_count.sum, kMinValue);

  // Check that the maximum private index that corresponds to a sum is decoded correctly.
  status = DecodePrivateIndexAsSumOrCount(MaxIndexForNumericValue(), metric_dimensions_, kMinValue,
                                          kMaxValue, kMaxCount, kNumIndexPoints, &event_vector,
                                          &sum_or_count);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({2, 300, 10}));
  EXPECT_EQ(sum_or_count.type, SumOrCount::SUM);
  EXPECT_EQ(sum_or_count.sum, kMaxValue);
}

TEST_F(PrivateIndexDecodingTest, DecodeAsSumOrCountValidCount) {
  std::vector<uint32_t> event_vector;
  SumOrCount sum_or_count;
  // Check that the minimum private index that corresponds to a count is decoded correctly.
  util::Status status =
      DecodePrivateIndexAsSumOrCount(MinIndexForCount(), metric_dimensions_, kMinValue, kMaxValue,
                                     kMaxCount, kNumIndexPoints, &event_vector, &sum_or_count);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({0, 0, 0}));
  EXPECT_EQ(sum_or_count.type, SumOrCount::COUNT);
  EXPECT_EQ(sum_or_count.count, 0u);

  // Check that the maximum private index that corresponds to a count is decoded correctly.
  status =
      DecodePrivateIndexAsSumOrCount(MaxIndexForCount(), metric_dimensions_, kMinValue, kMaxValue,
                                     kMaxCount, kNumIndexPoints, &event_vector, &sum_or_count);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({2, 300, 10}));
  EXPECT_EQ(sum_or_count.type, SumOrCount::COUNT);
  EXPECT_EQ(sum_or_count.count, kMaxCount);
}

TEST_F(PrivateIndexDecodingTest, DecodeAsSumOrCountInvalid) {
  std::vector<uint32_t> event_vector;
  SumOrCount sum_or_count;

  util::Status status = DecodePrivateIndexAsSumOrCount(
      MaxIndexForCount() + 1, metric_dimensions_, kMinValue, kMaxValue, kMaxCount, kNumIndexPoints,
      &event_vector, &sum_or_count);
  EXPECT_FALSE(status.ok());
}

TEST_F(PrivateIndexDecodingTest, DecodeAsHistogramBucketIndexValid) {
  std::vector<uint32_t> event_vector;
  uint32_t bucket_index;
  // Check that the minimum private index is decoded correctly.
  util::Status status = DecodePrivateIndexAsHistogramBucketIndex(
      0, metric_dimensions_, kMaxBucketIndex, &event_vector, &bucket_index);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({0, 0, 0}));
  EXPECT_EQ(bucket_index, 0u);

  // Check that the maximum private index is decoded correctly.
  status =
      DecodePrivateIndexAsHistogramBucketIndex(MaxIndexForHistogramBucket(), metric_dimensions_,
                                               kMaxBucketIndex, &event_vector, &bucket_index);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({2, 300, 10}));
  EXPECT_EQ(bucket_index, kMaxBucketIndex);
}

TEST_F(PrivateIndexDecodingTest, DecodeAsHistogramBucketIndexInvalid) {
  std::vector<uint32_t> event_vector;
  uint32_t bucket_index;

  util::Status status =
      DecodePrivateIndexAsHistogramBucketIndex(MaxIndexForHistogramBucket() + 1, metric_dimensions_,
                                               kMaxBucketIndex, &event_vector, &bucket_index);
  EXPECT_FALSE(status.ok());
}

TEST_F(PrivateIndexDecodingTest, DecodeAsHistogramBucketIndexAndCountValid) {
  std::vector<uint32_t> event_vector;
  uint32_t bucket_index;
  uint64_t bucket_count;
  // Check that the minimum private index is decoded correctly.
  util::Status status = DecodePrivateIndexAsHistogramBucketIndexAndCount(
      0, metric_dimensions_, kMaxBucketIndex, kMaxCount, kNumIndexPoints, &event_vector,
      &bucket_index, &bucket_count);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({0, 0, 0}));
  EXPECT_EQ(bucket_index, 0u);
  EXPECT_EQ(bucket_count, 0u);

  // Check that the maximum private index is decoded correctly.
  status = DecodePrivateIndexAsHistogramBucketIndexAndCount(
      MaxIndexForHistogramBucketAndCount(), metric_dimensions_, kMaxBucketIndex, kMaxCount,
      kNumIndexPoints, &event_vector, &bucket_index, &bucket_count);
  ASSERT_TRUE(status.ok());
  EXPECT_EQ(event_vector, std::vector<uint32_t>({2, 300, 10}));
  EXPECT_EQ(bucket_index, kMaxBucketIndex);
  EXPECT_EQ(bucket_count, kMaxCount);
}

TEST_F(PrivateIndexDecodingTest, DecodeAsHistogramBucketIndexAndCountInvalid) {
  std::vector<uint32_t> event_vector;
  uint32_t bucket_index;
  uint64_t bucket_count;

  util::Status status = DecodePrivateIndexAsHistogramBucketIndexAndCount(
      MaxIndexForHistogramBucketAndCount() + 1, metric_dimensions_, kMaxBucketIndex, kMaxCount,
      kNumIndexPoints, &event_vector, &bucket_index, &bucket_count);
  ASSERT_FALSE(status.ok());
}

}  // namespace cobalt