// Copyright 2024 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/metrics/structured/flushed_map.h"

#include <cstdint>
#include <optional>
#include <vector>

#include "base/files/file_path.h"
#include "base/files/file_util.h"
#include "base/files/scoped_temp_dir.h"
#include "base/run_loop.h"
#include "base/test/bind.h"
#include "base/test/task_environment.h"
#include "components/metrics/structured/lib/event_buffer.h"
#include "components/metrics/structured/proto/event_storage.pb.h"
#include "testing/gmock/include/gmock/gmock-matchers.h"
#include "testing/gtest/include/gtest/gtest.h"
#include "third_party/metrics_proto/structured_data.pb.h"

namespace metrics::structured {
namespace {

class TestEventBuffer : public EventBuffer<StructuredEventProto> {
 public:
  TestEventBuffer() : EventBuffer<StructuredEventProto>(ResourceInfo(1024)) {}

  // EventBuffer:
  Result AddEvent(StructuredEventProto event) override {
    events_.mutable_events()->Add(std::move(event));
    return Result::kOk;
  }

  void Purge() override { events_.Clear(); }

  google::protobuf::RepeatedPtrField<StructuredEventProto> Serialize()
      override {
    return events_.events();
  }

  void Flush(const base::FilePath& path, FlushedCallback callback) override {
    std::string content;
    EXPECT_TRUE(events_.SerializeToString(&content));
    EXPECT_TRUE(base::WriteFile(path, content));
    std::move(callback).Run(FlushedKey{
        .size = static_cast<int32_t>(content.size()),
        .path = path,
        .creation_time = base::Time::Now(),
    });
  }

  uint64_t Size() const override { return events_.events_size(); }

  const EventsProto& events() const { return events_; }

 private:
  EventsProto events_;
};

StructuredEventProto BuildTestEvent(int id) {
  StructuredEventProto event;
  event.set_device_project_id(id);
  return event;
}

TestEventBuffer BuildTestBuffer(const std::vector<int>& ids) {
  TestEventBuffer buffer;

  for (int id : ids) {
    EXPECT_EQ(buffer.AddEvent(BuildTestEvent(id)), Result::kOk);
  }

  return buffer;
}

EventsProto BuildTestEvents(const std::vector<int>& ids) {
  EventsProto events;

  for (int id : ids) {
    events.mutable_events()->Add(BuildTestEvent(id));
  }

  return events;
}

}  // namespace

class FlushedMapTest : public testing::Test {
 public:
  FlushedMapTest() = default;
  ~FlushedMapTest() override = default;

  void SetUp() override { ASSERT_TRUE(temp_dir_.CreateUniqueTempDir()); }

  base::FilePath GetDir() const {
    return temp_dir_.GetPath().Append(FILE_PATH_LITERAL("events"));
  }

  FlushedMap BuildFlushedMap(int32_t max_size = 2048) {
    return FlushedMap(GetDir(), max_size);
  }

  EventsProto ReadEventsProto(const base::FilePath& path) {
    std::string content;
    EXPECT_TRUE(base::ReadFileToString(path, &content));

    EventsProto events;
    EXPECT_TRUE(events.MergeFromString(content));
    return events;
  }

  void WriteToDisk(const base::FilePath& path, EventsProto&& events) {
    std::string content;
    EXPECT_TRUE(events.SerializeToString(&content));
    EXPECT_TRUE(base::WriteFile(path, content));
  }

  void Wait() { task_environment_.RunUntilIdle(); }

  FlushedKey BuildKey(const base::FilePath& path) {
    // This is a bit of a hack. We know what the key will be because we created
    // the file.
    FlushedKey key;
    key.path = path;
    // Manually create the size and other info for the key, since this is used
    // in size computations.
    base::File::Info info;
    EXPECT_TRUE(base::GetFileInfo(path, &info));
    key.size = info.size;
    key.creation_time = info.creation_time;
    return key;
  }

 private:
  base::test::TaskEnvironment task_environment_{
      base::test::TaskEnvironment::MainThreadType::UI,
      base::test::TaskEnvironment::ThreadPoolExecutionMode::QUEUED};

  base::ScopedTempDir temp_dir_;

 protected:
};

TEST_F(FlushedMapTest, FlushFile) {
  FlushedMap map = BuildFlushedMap();
  Wait();
  auto buffer = BuildTestBuffer({1, 2, 3});

  map.Flush(buffer,
            base::BindLambdaForTesting(
                [&](base::expected<FlushedKey, FlushError> key) {
                  EXPECT_TRUE(key.has_value());
                  EXPECT_TRUE(base::PathExists(base::FilePath(key->path)));
                }));
  Wait();

  const std::vector<FlushedKey>& keys = map.keys();
  EXPECT_EQ(keys.size(), 1ul);

  std::optional<int64_t> file_size = base::GetFileSize(keys[0].path);
  ASSERT_TRUE(file_size.has_value());
  EXPECT_EQ(keys[0].size, static_cast<int32_t>(file_size.value()));
}

TEST_F(FlushedMapTest, ReadFile) {
  FlushedMap map = BuildFlushedMap();
  Wait();

  auto buffer = BuildTestBuffer({1, 2, 3});
  const EventsProto& events = buffer.events();

  map.Flush(buffer,
            base::BindLambdaForTesting(
                [&](base::expected<FlushedKey, FlushError> key) {
                  EXPECT_TRUE(key.has_value());
                  EXPECT_TRUE(base::PathExists(base::FilePath(key->path)));
                }));
  Wait();

  auto key = map.keys().front();

  std::optional<EventsProto> read_events = map.ReadKey(key);
  EXPECT_EQ(read_events->events_size(), events.events_size());

  for (int i = 0; i < read_events->events_size(); ++i) {
    EXPECT_EQ(read_events->events(i).device_project_id(),
              events.events(i).device_project_id());
  }
}

TEST_F(FlushedMapTest, ReadKeyMissingFile) {
  FlushedKey key;
  key.path = GetDir().Append(FILE_PATH_LITERAL("nonexistent"));
  EXPECT_EQ(FlushedMap::ReadKey(key), std::nullopt);
}

TEST_F(FlushedMapTest, UniqueFlushes) {
  FlushedMap map = BuildFlushedMap();
  Wait();

  auto events = BuildTestBuffer({1, 2, 3});
  map.Flush(events,
            base::BindLambdaForTesting(
                [&](base::expected<FlushedKey, FlushError> key) {
                  EXPECT_TRUE(key.has_value());
                  EXPECT_TRUE(base::PathExists(base::FilePath(key->path)));
                }));
  Wait();

  auto events2 = BuildTestBuffer({4, 5, 6});
  map.Flush(events2,
            base::BindLambdaForTesting(
                [&](base::expected<FlushedKey, FlushError> key) {
                  EXPECT_TRUE(key.has_value());
                  EXPECT_TRUE(base::PathExists(base::FilePath(key->path)));
                }));
  Wait();

  EXPECT_EQ(map.keys().size(), 2ul);
  const auto& key1 = map.keys()[0];
  const auto& key2 = map.keys()[1];

  EXPECT_NE(key1.path, key2.path);
}

TEST_F(FlushedMapTest, DeleteKey) {
  FlushedMap map = BuildFlushedMap();
  Wait();

  auto events = BuildTestBuffer({1, 2, 3});
  map.Flush(events, base::BindLambdaForTesting(
                        [&](base::expected<FlushedKey, FlushError> key) {
                          EXPECT_TRUE(key.has_value());
                          EXPECT_TRUE(base::PathExists(key->path));
                        }));
  Wait();

  const std::vector<FlushedKey>& keys = map.keys();
  auto key = map.keys().front();
  EXPECT_EQ(keys.size(), 1ul);
  EXPECT_TRUE(base::PathExists(key.path));

  map.DeleteKey(key);
  Wait();

  EXPECT_FALSE(base::PathExists(key.path));
}

TEST_F(FlushedMapTest, LoadPreviousSessionKeys) {
  EXPECT_TRUE(base::CreateDirectory(GetDir()));
  auto events = BuildTestEvents({1, 2, 3});
  auto events2 = BuildTestEvents({4, 5, 6});

  base::FilePath path1 = GetDir().Append(FILE_PATH_LITERAL("events"));
  base::FilePath path2 = GetDir().Append(FILE_PATH_LITERAL("events2"));

  WriteToDisk(path1, std::move(events));
  // Force a small difference in the creation time of the two files.
  base::PlatformThreadBase::Sleep(base::Seconds(1));
  WriteToDisk(path2, std::move(events2));

  FlushedMap map = BuildFlushedMap();
  Wait();

  const std::vector<FlushedKey>& keys = map.keys();
  EXPECT_EQ(keys.size(), 2ul);

  // The order the files are loaded in is unknown.
  std::vector<base::FilePath> paths;
  for (const auto& key : keys) {
    paths.push_back(base::FilePath(key.path));
  }

  EXPECT_THAT(paths, testing::ElementsAre(path1, path2));
}

TEST_F(FlushedMapTest, ExceedQuota) {
  FlushedMap map = BuildFlushedMap(/*max_size=*/64);
  Wait();

  auto events = BuildTestBuffer({1, 2, 3});
  map.Flush(events, base::BindLambdaForTesting(
                        [&](base::expected<FlushedKey, FlushError> key) {
                          EXPECT_TRUE(key.has_value());
                          EXPECT_TRUE(base::PathExists(key->path));
                        }));
  Wait();

  auto events2 = BuildTestBuffer({1, 2, 3});
  map.Flush(events2, base::BindLambdaForTesting(
                         [&](base::expected<FlushedKey, FlushError> key) {
                           EXPECT_FALSE(key.has_value());
                           const FlushError err = key.error();
                           EXPECT_EQ(err, kQuotaExceeded);
                         }));
  Wait();
}

TEST_F(FlushedMapTest, DeferPurgeBeforeInit) {
  EXPECT_TRUE(base::CreateDirectory(GetDir()));
  auto events = BuildTestEvents({1, 2, 3});
  base::FilePath path = GetDir().Append(FILE_PATH_LITERAL("events"));
  WriteToDisk(path, std::move(events));

  FlushedMap map = BuildFlushedMap();
  map.Purge();
  Wait();
  EXPECT_TRUE(map.keys().empty());
  EXPECT_FALSE(base::PathExists(path));
}

TEST_F(FlushedMapTest, DeferDeleteKeyBeforeInit) {
  EXPECT_TRUE(base::CreateDirectory(GetDir()));
  auto events = BuildTestEvents({1, 2, 3});
  base::FilePath path = GetDir().Append(FILE_PATH_LITERAL("events"));
  WriteToDisk(path, std::move(events));

  FlushedMap map = BuildFlushedMap();
  map.DeleteKey(BuildKey(path));
  Wait();
  EXPECT_TRUE(map.keys().empty());
  EXPECT_FALSE(base::PathExists(path));
}

TEST_F(FlushedMapTest, DeferFlushBeforeInit) {
  // Manually create the directory, because the flush happens before the
  // FlushedMap has a chance to create it on its background thread.
  ASSERT_TRUE(base::CreateDirectory(GetDir()));
  FlushedMap map = BuildFlushedMap();
  auto buffer = BuildTestBuffer({1, 2, 3});
  bool callback_ran = false;
  map.Flush(buffer,
            base::BindLambdaForTesting(
                [&](base::expected<FlushedKey, FlushError> key) {
                  EXPECT_TRUE(key.has_value());
                  EXPECT_TRUE(base::PathExists(base::FilePath(key->path)));
                  callback_ran = true;
                }));
  Wait();
  EXPECT_TRUE(callback_ran);
  EXPECT_EQ(map.keys().size(), 1ul);
}

TEST_F(FlushedMapTest, DeleteKeysWithMissingKeys) {
  FlushedMap map = BuildFlushedMap();
  Wait();

  auto buffer1 = BuildTestBuffer({1});
  map.Flush(buffer1, base::BindLambdaForTesting(
                         [&](base::expected<FlushedKey, FlushError> key) {}));
  auto buffer2 = BuildTestBuffer({2});
  base::FilePath path2;
  map.Flush(buffer2, base::BindLambdaForTesting(
                         [&](base::expected<FlushedKey, FlushError> key) {
                           path2 = key->path;
                         }));
  Wait();

  ASSERT_EQ(map.keys().size(), 2ul);

  FlushedKey missing_key;
  missing_key.path = GetDir().Append(FILE_PATH_LITERAL("missing"));

  std::vector<FlushedKey> keys_to_delete = {map.keys()[0], missing_key};
  map.DeleteKeys(keys_to_delete);
  Wait();

  EXPECT_EQ(map.keys().size(), 1ul);
  EXPECT_EQ(map.keys()[0].path, path2);
}

TEST_F(FlushedMapTest, DeferMultipleOperations) {
  // Create two files on disk.
  ASSERT_TRUE(base::CreateDirectory(GetDir()));
  auto events1 = BuildTestEvents({1, 2, 3});
  base::FilePath path1 = GetDir().Append(FILE_PATH_LITERAL("events1"));
  WriteToDisk(path1, std::move(events1));
  auto events2 = BuildTestEvents({4, 5, 6});
  base::FilePath path2 = GetDir().Append(FILE_PATH_LITERAL("events2"));
  WriteToDisk(path2, std::move(events2));

  FlushedMap map = BuildFlushedMap();

  // Before initialization, delete one of the keys and flush a new buffer.
  map.DeleteKey(BuildKey(path1));

  auto buffer = BuildTestBuffer({7, 8, 9});
  map.Flush(buffer, base::BindLambdaForTesting(
                        [&](base::expected<FlushedKey, FlushError> key) {}));

  Wait();

  // After initialization, we should have two keys: the one that wasn't deleted
  // and the newly flushed one.
  EXPECT_EQ(map.keys().size(), 2ul);

  std::vector<base::FilePath> paths;
  for (const auto& key : map.keys()) {
    paths.push_back(key.path);
  }

  // The order is not guaranteed.
  EXPECT_THAT(paths, testing::UnorderedElementsAre(
                         path2, testing::Not(testing::Eq(path1))));
}

}  // namespace metrics::structured