// Copyright 2020 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/local_aggregation_1_1/local_aggregate_storage/immediate_local_aggregate_storage.h"

#include <memory>

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

#include "absl/strings/escaping.h"
#include "absl/strings/str_cat.h"
#include "src/lib/statusor/statusor.h"
#include "src/lib/util/testing/test_with_files.h"
#include "src/local_aggregation_1_1/local_aggregate_storage/local_aggregate_storage.h"
#include "src/local_aggregation_1_1/testing/test_registry.cb.h"
#include "src/logger/project_context_factory.h"
#include "src/logging.h"
#include "src/registry/cobalt_registry.pb.h"

namespace cobalt::local_aggregation {
using lib::statusor::StatusOr;
using ::testing::Contains;
using ::testing::Not;
using MetricAggregateRef = LocalAggregateStorage::MetricAggregateRef;

namespace {

std::unique_ptr<CobaltRegistry> GetRegistry() {
  std::string bytes;
  if (!absl::Base64Unescape(kCobaltRegistryBase64, &bytes)) {
    LOG(ERROR) << "Unable to decode Base64 String";
    return nullptr;
  }

  auto registry = std::make_unique<CobaltRegistry>();
  if (!registry->ParseFromString(bytes)) {
    LOG(ERROR) << "Unable to parse registry from bytes";
    return nullptr;
  }

  return registry;
}

}  // namespace

class ImmediateLocalAggregateStorageTest : public util::testing::TestWithFiles {
 private:
  void SetUp() override {
    MakeTestFolder();
    ReplaceRegistry();
  }

 public:
  void ReplaceRegistry(std::unique_ptr<CobaltRegistry> registry = GetRegistry()) {
    global_project_context_factory_ =
        std::make_unique<logger::ProjectContextFactory>(std::move(registry));
    storage_ = std::make_unique<ImmediateLocalAggregateStorage>(
        test_folder(), fs(), global_project_context_factory_.get());
  }

 protected:
  std::unique_ptr<logger::ProjectContextFactory> global_project_context_factory_;
  std::unique_ptr<ImmediateLocalAggregateStorage> storage_;
};

TEST_F(ImmediateLocalAggregateStorageTest, MakesExpectedFiles) {
  // Root directory should contain expected customer_id: 123
  ASSERT_THAT(fs()->ListFiles(test_folder()).ConsumeValueOrDie(), Contains("123"));

  // Customer 123 directory should contain expected project: 100
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123")).ConsumeValueOrDie(),
              Contains("100"));

  // Customer 123 Project 100 directory should *not* contain the metric file: 1
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123/100")).ConsumeValueOrDie(),
              Not(Contains("1")));

  StatusOr<MetricAggregateRef> agg_or = storage_->GetMetricAggregate(123, 100, 1);
  ASSERT_TRUE(agg_or.ok());
  agg_or.ConsumeValueOrDie().Save();

  // Customer 123 Project 100 directory should contain the metric file: 1
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123/100")).ConsumeValueOrDie(),
              Contains("1"));
}

TEST_F(ImmediateLocalAggregateStorageTest, CanReadAlreadyWrittenFiles) {
  {
    StatusOr<MetricAggregateRef> agg_or = storage_->GetMetricAggregate(123, 100, 1);
    ASSERT_TRUE(agg_or.ok());

    MetricAggregateRef agg = agg_or.ConsumeValueOrDie();
    agg.aggregate()->set_version(100);
    ASSERT_TRUE(agg.Save().ok());
  }

  ReplaceRegistry();

  StatusOr<MetricAggregateRef> agg_or = storage_->GetMetricAggregate(123, 100, 1);
  ASSERT_TRUE(agg_or.ok());
  ASSERT_TRUE(agg_or.ConsumeValueOrDie().Save().ok());
}

TEST_F(ImmediateLocalAggregateStorageTest, CleansUpOldMetrics) {
  StatusOr<MetricAggregateRef> agg_or = storage_->GetMetricAggregate(123, 100, 1);
  ASSERT_TRUE(agg_or.ok());
  agg_or.ConsumeValueOrDie().Save();

  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123/100")).ConsumeValueOrDie(),
              Contains("1"));

  auto registry = GetRegistry();
  google::protobuf::RepeatedPtrField<MetricDefinition>* metrics =
      registry->mutable_customers(0)->mutable_projects(0)->mutable_metrics();
  metrics->erase(metrics->begin() + kOccurrenceMetricMetricIndex);
  ReplaceRegistry(std::move(registry));

  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123/100")).ConsumeValueOrDie(),
              Not(Contains("1")));
}

TEST_F(ImmediateLocalAggregateStorageTest, CleansUpOldProjects) {
  StatusOr<MetricAggregateRef> agg_or = storage_->GetMetricAggregate(123, 100, 1);
  ASSERT_TRUE(agg_or.ok());
  agg_or.ConsumeValueOrDie().Save();

  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123")).ConsumeValueOrDie(),
              Contains("100"));
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123/100")).ConsumeValueOrDie(),
              Contains("1"));

  auto registry = GetRegistry();
  registry->mutable_customers(0)->mutable_projects()->RemoveLast();
  ReplaceRegistry(std::move(registry));

  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123")).ConsumeValueOrDie(),
              Not(Contains("100")));
}

TEST_F(ImmediateLocalAggregateStorageTest, DeleteDataWorks) {
  StatusOr<MetricAggregateRef> agg_or = storage_->GetMetricAggregate(123, 100, 1);
  ASSERT_TRUE(agg_or.ok());
  agg_or.ConsumeValueOrDie().Save();

  ASSERT_THAT(fs()->ListFiles(test_folder()).ConsumeValueOrDie(), Contains("123"));
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123")).ConsumeValueOrDie(),
              Contains("100"));
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_folder(), "/123/100")).ConsumeValueOrDie(),
              Contains("1"));

  storage_->DeleteData();

  ASSERT_THAT(fs()->ListFiles(test_folder()).ConsumeValueOrDie(), Not(Contains("123")));
}

class ImmediateLocalAggregateStorageMissingDirectoryTest : public util::testing::TestWithFiles {
 private:
  void SetUp() override {
    MakeTestFolder();
    ReplaceRegistry();
  }

 public:
  void ReplaceRegistry(std::unique_ptr<CobaltRegistry> registry = GetRegistry()) {
    global_project_context_factory_ =
        std::make_unique<logger::ProjectContextFactory>(std::move(registry));
    storage_ = std::make_unique<ImmediateLocalAggregateStorage>(
        test_subfolder(), fs(), global_project_context_factory_.get());
  }

 protected:
  std::string test_subfolder() { return absl::StrCat(test_folder(), "/sub"); }

  std::unique_ptr<logger::ProjectContextFactory> global_project_context_factory_;
  std::unique_ptr<ImmediateLocalAggregateStorage> storage_;
};

TEST_F(ImmediateLocalAggregateStorageMissingDirectoryTest, DontAssumeDirectoryExists) {
  // Root directory should contain expected customer_id: 123
  ASSERT_THAT(fs()->ListFiles(test_subfolder()).ConsumeValueOrDie(), Contains("123"));

  // Customer 123 directory should contain expected project: 100
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_subfolder(), "/123")).ConsumeValueOrDie(),
              Contains("100"));

  // Customer 123 Project 100 directory should *not* contain the metric file: 1
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_subfolder(), "/123/100")).ConsumeValueOrDie(),
              Not(Contains("1")));

  StatusOr<MetricAggregateRef> agg_or = storage_->GetMetricAggregate(123, 100, 1);
  ASSERT_TRUE(agg_or.ok());
  agg_or.ConsumeValueOrDie().Save();

  // Customer 123 Project 100 directory should contain the metric file: 1
  ASSERT_THAT(fs()->ListFiles(absl::StrCat(test_subfolder(), "/123/100")).ConsumeValueOrDie(),
              Contains("1"));
}

}  // namespace cobalt::local_aggregation