// 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/backfill_manager.h"

#include <memory>

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

#include "src/lib/util/clock.h"
#include "src/lib/util/datetime_util.h"
#include "src/local_aggregation_1_1/aggregation_procedures/aggregation_procedure.h"
#include "src/public/lib/statusor/status_macros.h"

namespace cobalt::local_aggregation {

using TimeInfo = util::TimeInfo;

namespace {
MetricDefinition MetricWithTimeZonePolicy(MetricDefinition::TimeZonePolicy policy) {
  MetricDefinition metric;
  metric.set_time_zone_policy(policy);
  return metric;
}

std::vector<uint32_t> ToHourIds(std::vector<TimeInfo> times) {
  std::vector<uint32_t> output;
  std::transform(times.begin(), times.end(), std::back_inserter(output),
                 [](TimeInfo time) { return time.hour_id; });
  return output;
}

std::vector<uint32_t> ToDayIndices(std::vector<TimeInfo> times) {
  std::vector<uint32_t> output;
  std::transform(times.begin(), times.end(), std::back_inserter(output),
                 [](TimeInfo time) { return time.day_index; });
  return output;
}
}  // namespace

TEST(BackfillManager, HourlyBackfill) {
  const uint32_t kNumHoursToGenerate = 5;
  BackfillManager manager(2);
  util::UtcTimeConverter converter;
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::UTC);

  std::chrono::system_clock::time_point start_time(std::chrono::hours(10001) /
                                                   util::kHourIdsPerHour);
  std::chrono::system_clock::time_point end_time =
      start_time + std::chrono::hours(kNumHoursToGenerate + 1);

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo last_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  ASSERT_EQ(last_time_info.hour_id, 10001);

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(last_time_info, end_time, converter, metric, false));
  EXPECT_EQ(ToHourIds(backfill), std::vector<uint32_t>({10003, 10005, 10007, 10009, 10011}));
}

TEST(BackfillManager, HourlyBackfillEdgecase) {
  BackfillManager manager(2);
  util::UtcTimeConverter converter;

  // This edgecase was discovered in testing.
  // When end_epoch is set to 1 more than the day index corresponding to `end_time`, this will fail.
  // When end_epoch is set to 2 more than the day index corresponding to `end_time`, this passes.
  TimeInfo start = TimeInfo::FromHourId(0);
  std::chrono::system_clock::time_point end_time(
      std::chrono::hours(896489 / util::kHourIdsPerHour + 1));

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(start, end_time, converter,
                                MetricWithTimeZonePolicy(MetricDefinition::UTC), false));

  std::vector<uint32_t> hour_sequence = ToHourIds(backfill);
  EXPECT_EQ(hour_sequence[hour_sequence.size() - 1], 896489);
}

TEST(BackfillManager, HourlyBackstopDayAligned) {
  BackfillManager manager(1);
  util::UtcTimeConverter converter;
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::UTC);

  const uint32_t kStartDayIndex = 100;
  const uint32_t kStartHourId = (kStartDayIndex * util::kHourIdsPerDay) + 1;
  // Set the end time to the first hour of the next day.
  std::chrono::system_clock::time_point end_time(
      std::chrono::hours(util::kNumHoursPerDay * (kStartDayIndex + 2)));

  CB_ASSERT_OK_AND_ASSIGN(std::vector<TimeInfo> backfill,
                          manager.CalculateBackfill(TimeInfo::FromHourId(kStartHourId), end_time,
                                                    converter, metric, false));

  std::vector<uint32_t> hour_sequence = ToHourIds(backfill);
  ASSERT_EQ(hour_sequence.size(), 24);
  EXPECT_EQ(*hour_sequence.rbegin(), util::TimePointToHourIdUtc(end_time) - 2);
}

TEST(BackfillManager, HourlyBackstopDayUnaligned) {
  BackfillManager manager(1);
  util::UtcTimeConverter converter;

  const uint32_t kStartDayIndex = 100;
  const uint32_t kStartHourId = (kStartDayIndex * util::kHourIdsPerDay) + 1;
  // Set the end time to 3 hours before the end of the same day.
  std::chrono::system_clock::time_point end_time(
      std::chrono::hours(util::kNumHoursPerDay * (kStartDayIndex + 1) - 3));

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(TimeInfo::FromHourId(kStartHourId), end_time, converter,
                                MetricWithTimeZonePolicy(MetricDefinition::UTC), false));

  std::vector<uint32_t> hour_sequence = ToHourIds(backfill);
  ASSERT_EQ(*hour_sequence.rbegin(), util::TimePointToHourIdUtc(end_time) - 2);
  EXPECT_EQ(hour_sequence.size(), 20);
}

TEST(BackfillManager, DailyBackfill) {
  BackfillManager manager(5);
  util::UtcTimeConverter converter;

  const uint32_t kStartDayIndex = 10001;
  const uint32_t kNumDaysToGenerate = 5;
  std::chrono::system_clock::time_point end_time(
      std::chrono::hours(util::kNumHoursPerDay * (kStartDayIndex + kNumDaysToGenerate + 1)));

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(TimeInfo::FromDayIndex(kStartDayIndex), end_time, converter,
                                MetricWithTimeZonePolicy(MetricDefinition::UTC), true));
  EXPECT_EQ(ToDayIndices(backfill), std::vector<uint32_t>({10002, 10003, 10004, 10005, 10006}));
}

TEST(BackfillManager, DailyBackstop) {
  BackfillManager manager(10);
  util::UtcTimeConverter converter;

  const uint32_t kStartDayIndex = 100;
  const uint32_t kLastDayIndex = 120;
  std::chrono::system_clock::time_point end_time(
      std::chrono::hours(util::kNumHoursPerDay * (kLastDayIndex + 1)));

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(TimeInfo::FromDayIndex(kStartDayIndex), end_time, converter,
                                MetricWithTimeZonePolicy(MetricDefinition::UTC), true));

  std::vector<uint32_t> day_sequence = ToDayIndices(backfill);
  ASSERT_EQ(day_sequence.size(), 10);
  EXPECT_EQ(*day_sequence.rbegin(), kLastDayIndex);
}

// Test backfill calculation for an OTHER_TIME_ZONE metric, with a time zone that agrees with UTC.
// The result should be the same as for a UTC metric.
TEST(BackfillManager, HourlyBackfillOtherTimeZone) {
  const uint32_t kNumHoursToGenerate = 5;
  BackfillManager manager(2);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  std::chrono::system_clock::time_point start_time(
      std::chrono::hours(10001 / util::kHourIdsPerHour));
  std::chrono::system_clock::time_point end_time =
      start_time + std::chrono::hours(kNumHoursToGenerate + 1);

  util::FakeCivilTimeConverter converter(/*start_utc_offset=*/0, /*start_isdst=*/false);

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo start_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  ASSERT_EQ(start_time_info.hour_id, 10001);

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(start_time_info, end_time, converter, metric, false));
  EXPECT_EQ(ToHourIds(backfill), std::vector<uint32_t>({10003, 10005, 10007, 10009, 10011}));
}

// Test backfill calculation for an OTHER_TIME_ZONE metric with a time zone at an offset from UTC.
TEST(BackfillManager, HourlyBackfillOtherTimeZoneUtcOffset) {
  const uint32_t kNumHoursToGenerate = 5;
  BackfillManager manager(2);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  std::chrono::system_clock::time_point start_time(
      std::chrono::hours(10001 / util::kHourIdsPerHour));
  std::chrono::system_clock::time_point end_time =
      start_time + std::chrono::hours(kNumHoursToGenerate + 1);

  util::FakeCivilTimeConverter converter(/*start_utc_offset=*/8, /*start_isdst=*/false);

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo start_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  ASSERT_EQ(start_time_info.hour_id, 10017);

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(start_time_info, end_time, converter, metric, false));
  EXPECT_EQ(ToHourIds(backfill), std::vector<uint32_t>({10019, 10021, 10023, 10025, 10027}));
}

// Test backfill calculation for an OTHER_TIME_ZONE metric, in a case where the backstop time falls
// multiple hours after the beginning of the backfill period in the metric's time zone.
TEST(BackfillManager, HourlyBackfillOtherTimeZoneEarlyBackstop) {
  const uint32_t kNumHoursToGenerate = 24;
  BackfillManager manager(1);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  std::chrono::system_clock::time_point start_time(std::chrono::hours(util::kNumHoursPerDay * 100));
  std::chrono::system_clock::time_point end_time =
      start_time + std::chrono::hours(kNumHoursToGenerate + 1);

  // The end time has day index 101 with respect to UTC, and the backfill period is 2 days.
  // The backfill manager should only return hour IDs that correspond to day indices 100 and 101.
  ASSERT_EQ(util::TimePointToDayIndexUtc(end_time), 101);

  util::FakeCivilTimeConverter converter(/*start_utc_offset=*/-23, /*start_isdst=*/false);

  // All times between `start_time` and `end_time`-3h have day index 99 with respect to `metric`'s
  // time zone. Only the hours `end_time`-2h and `end_time`-1h (hour IDs 4801 and 4803) fall within
  // the backfill period.
  CB_ASSERT_OK_AND_ASSIGN(TimeInfo start_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  ASSERT_EQ(start_time_info.hour_id, 4755);
  ASSERT_EQ(start_time_info.day_index, 99);

  CB_ASSERT_OK_AND_ASSIGN(
      TimeInfo boundary_info,
      TimeInfo::FromTimePoint(end_time - std::chrono::hours(3), converter, metric));
  ASSERT_EQ(boundary_info.hour_id, 4799);
  ASSERT_EQ(boundary_info.day_index, 99);

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo end_time_info,
                          TimeInfo::FromTimePoint(end_time, converter, metric));
  ASSERT_EQ(end_time_info.hour_id, 4805);
  ASSERT_EQ(end_time_info.day_index, 100);

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(start_time_info, end_time, converter, metric, false));
  EXPECT_EQ(ToHourIds(backfill), std::vector<uint32_t>({4801, 4803}));
}

// Test backfill calculation for an OTHER_TIME_ZONE metric, with a time zone that agrees with UTC
// except that the DST flag is set.
TEST(BackfillManager, HourlyBackfillOtherTimeZoneDst) {
  const uint32_t kNumHoursToGenerate = 5;
  BackfillManager manager(2);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  std::chrono::system_clock::time_point start_time(
      std::chrono::hours(10001 / util::kHourIdsPerHour));
  std::chrono::system_clock::time_point end_time =
      start_time + std::chrono::hours(kNumHoursToGenerate + 1);

  util::FakeCivilTimeConverter converter(/*start_utc_offset=*/0, /*start_isdst=*/true);

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo start_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  ASSERT_EQ(start_time_info.hour_id, 10000);

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(start_time_info, end_time, converter, metric, false));
  EXPECT_EQ(ToHourIds(backfill), std::vector<uint32_t>({10002, 10004, 10006, 10008, 10010}));
}

// Test backfill calculation when DST starts during the backfill period.
TEST(BackfillManager, HourlyBackfillWithDstStart) {
  const uint32_t kNumHoursToGenerate = 5;
  BackfillManager manager(2);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  std::chrono::system_clock::time_point start_time(
      std::chrono::hours(10001 / util::kHourIdsPerHour));
  std::chrono::system_clock::time_point end_time =
      start_time + std::chrono::hours(kNumHoursToGenerate + 1);

  // Use a civil time converter that simulates a transition from GMT (UTC+0, not DST) to British
  // Summer Time (UTC+1, DST), with the transition occurring 2 hours after the last observation was
  // generated.
  util::FakeCivilTimeConverter converter(
      /*start_utc_offset=*/0, /*start_isdst=*/false,
      /*end_utc_offset=*/1, /*end_isdst=*/true,
      /*threshold=*/start_time + 2 * util::kOneHour);

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo start_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  ASSERT_EQ(start_time_info.hour_id, 10001);

  CB_ASSERT_OK_AND_ASSIGN(std::vector<TimeInfo> backfill,
                          manager.CalculateBackfill(start_time_info, end_time, converter, metric,
                                                    /*is_daily=*/false));
  // The first hour ID in the backfill period is odd (not DST) and the remainder are even (DST).
  EXPECT_EQ(ToHourIds(backfill), std::vector<uint32_t>({10003, 10006, 10008, 10010, 10012}));
}

// Test backfill calculation when DST ends during the backfill period.
TEST(BackfillManager, HourlyBackfillWithDstEnd) {
  const uint32_t kNumHoursToGenerate = 5;
  BackfillManager manager(2);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  std::chrono::system_clock::time_point start_time(
      std::chrono::hours(10001 / util::kHourIdsPerHour));
  std::chrono::system_clock::time_point end_time =
      start_time + std::chrono::hours(kNumHoursToGenerate + 1);

  // Use a civil time converter that simulates a transition from British Summer Time (UTC+1, DST) to
  // GMT (UTC+0, not DST) with the transition occurring 2 hours after the last observation was
  // generated.
  util::FakeCivilTimeConverter converter(
      /*start_utc_offset=*/1, /*start_isdst=*/true,
      /*end_utc_offset=*/0, /*end_isdst=*/false,
      /*threshold=*/start_time + std::chrono::hours(2));

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo start_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  ASSERT_EQ(start_time_info.hour_id, 10002);

  CB_ASSERT_OK_AND_ASSIGN(std::vector<TimeInfo> backfill,
                          manager.CalculateBackfill(start_time_info, end_time, converter, metric,
                                                    /*is_daily=*/false));
  // The first hour ID in the backfill period is even (DST) and the remainder are odd (not DST).
  EXPECT_EQ(ToHourIds(backfill), std::vector<uint32_t>({10004, 10005, 10007, 10009, 10011}));
}

// Tests that the sequence of hour IDs returned by the BackfillManager agrees with the sequence of
// hour IDs assigned to logged events, when the time period includes the start of DST.
TEST(BackfillManager, HourlyBackfillSameAsLoggedHoursWithDstStart) {
  BackfillManager manager(2);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  std::chrono::system_clock::time_point start_time(
      std::chrono::hours(10001 / util::kHourIdsPerHour));
  std::chrono::system_clock::time_point end_time = start_time + util::kOneDay;

  // Use a civil time converter that simulates a transition from GMT (UTC+0, not DST) to British
  // Summer Time (UTC+1, DST), with the transition occurring 2 hours after the last observation was
  // generated.
  util::FakeCivilTimeConverter converter(
      /*start_utc_offset=*/0, /*start_isdst=*/false,
      /*end_utc_offset=*/1, /*end_isdst=*/true,
      /*threshold=*/start_time + 2 * util::kOneHour);

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo last_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  CB_ASSERT_OK_AND_ASSIGN(std::vector<TimeInfo> backfill,
                          manager.CalculateBackfill(last_time_info, end_time, converter, metric,
                                                    /*is_daily=*/false));

  // Generate all possible hour IDs from the start to end of the time window.
  std::set<uint32_t> hour_ids;
  std::chrono::system_clock::time_point time = start_time;
  while (time < end_time) {
    CB_ASSERT_OK_AND_ASSIGN(uint32_t hour_id, util::TimePointToHourId(time, converter, metric));
    hour_ids.insert(hour_id);
    time += std::chrono::minutes(1);
  }
  // The hour ID of the last generated observations shouldn't show up in `backfill`.
  hour_ids.erase(last_time_info.hour_id);
  EXPECT_THAT(ToHourIds(backfill), testing::UnorderedElementsAreArray(hour_ids));
}

// Tests that the sequence of hour IDs returned by the BackfillManager agrees with the sequence of
// hour IDs assigned to logged events, when the time period includes the end of DST.
TEST(BackfillManager, HourlyBackfillSameAsLoggedHoursWithDstEnd) {
  BackfillManager manager(2);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  std::chrono::system_clock::time_point start_time(
      std::chrono::hours(10001 / util::kHourIdsPerHour));
  std::chrono::system_clock::time_point end_time = start_time + util::kOneDay;

  // Use a civil time converter that simulates a transition from British Summer Time (UTC+1, DST) to
  // GMT (UTC+0, not DST) with the transition occurring 2 hours after the last observation was
  // generated.
  util::FakeCivilTimeConverter converter(
      /*start_utc_offset=*/1, /*start_isdst=*/true,
      /*end_utc_offset=*/0, /*end_isdst=*/false,
      /*threshold=*/start_time + std::chrono::hours(2));

  CB_ASSERT_OK_AND_ASSIGN(TimeInfo last_time_info,
                          TimeInfo::FromTimePoint(start_time, converter, metric));
  CB_ASSERT_OK_AND_ASSIGN(std::vector<TimeInfo> backfill,
                          manager.CalculateBackfill(last_time_info, end_time, converter, metric,
                                                    /*is_daily=*/false));

  // Generate all possible hour IDs from the start to end of the time window.
  std::set<uint32_t> hour_ids;
  std::chrono::system_clock::time_point time = start_time;
  while (time < end_time) {
    CB_ASSERT_OK_AND_ASSIGN(uint32_t hour_id, util::TimePointToHourId(time, converter, metric));
    hour_ids.insert(hour_id);
    time += std::chrono::minutes(1);
  }
  // The hour ID of the last generated observations shouldn't show up in `backfill`.
  hour_ids.erase(last_time_info.hour_id);
  EXPECT_THAT(ToHourIds(backfill), testing::UnorderedElementsAreArray(hour_ids));
}

// Test a simple case of daily backfill for an OTHER_TIME_ZONE metric.
TEST(BackfillManager, DailyBackfillOtherTimeZone) {
  BackfillManager manager(5);
  util::FakeCivilTimeConverter converter(/*start_utc_offset=*/8, /*start_isdst=*/false);

  const uint32_t kStartDayIndex = 10001;
  const uint32_t kNumDaysToGenerate = 5;
  std::chrono::system_clock::time_point end_time(
      std::chrono::hours(util::kNumHoursPerDay * (kStartDayIndex + kNumDaysToGenerate + 1)));

  CB_ASSERT_OK_AND_ASSIGN(
      std::vector<TimeInfo> backfill,
      manager.CalculateBackfill(TimeInfo::FromDayIndex(kStartDayIndex), end_time, converter,
                                MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE), true));
  EXPECT_EQ(ToDayIndices(backfill), std::vector<uint32_t>({10002, 10003, 10004, 10005, 10006}));
}

// Test daily backfill when the final day index in the metric's time zone is smaller than the final
// day index in UTC.
TEST(BackfillManager, DailyBackfillShiftedOtherTimeZone) {
  BackfillManager manager(5);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  util::FakeCivilTimeConverter converter(/*start_utc_offset=*/-23, /*start_isdst=*/false);

  const uint32_t kStartDayIndex = 10001;
  const uint32_t kNumDaysToGenerate = 5;
  std::chrono::system_clock::time_point end_time(
      std::chrono::hours(util::kNumHoursPerDay * (kStartDayIndex + kNumDaysToGenerate + 1)));

  // The day index of `end_time` in `metric`'s time zone is 10006 (rather than 10007 in UTC).
  CB_ASSERT_OK_AND_ASSIGN(uint32_t end_day_index,
                          util::TimePointToDayIndex(end_time, converter, metric));
  ASSERT_EQ(end_day_index, 10006);

  CB_ASSERT_OK_AND_ASSIGN(std::vector<TimeInfo> backfill,
                          manager.CalculateBackfill(TimeInfo::FromDayIndex(kStartDayIndex),
                                                    end_time, converter, metric, true));
  // Since the day index of `end_time` is 10006, only return 4 day indices (instead of 5 as would be
  // expected for UTC).
  EXPECT_EQ(ToDayIndices(backfill), std::vector<uint32_t>({10002, 10003, 10004, 10005}));
}

// Test daily backfill when the backstop day is multiple days later than the last day for which
// observations were generated.
TEST(BackfillManager, DailyBackstopOtherTimeZone) {
  BackfillManager manager(2);
  MetricDefinition metric = MetricWithTimeZonePolicy(MetricDefinition::OTHER_TIME_ZONE);

  util::FakeCivilTimeConverter converter(/*start_utc_offset=*/-23, /*start_isdst=*/false);

  const uint32_t kStartDayIndex = 10001;
  const uint32_t kNumDaysToGenerate = 5;
  std::chrono::system_clock::time_point end_time(
      std::chrono::hours(util::kNumHoursPerDay * (kStartDayIndex + kNumDaysToGenerate + 1)));

  // The day index of `end_time` in `metric`'s time zone is 10006.
  CB_ASSERT_OK_AND_ASSIGN(uint32_t end_day_index,
                          util::TimePointToDayIndex(end_time, converter, metric));
  ASSERT_EQ(end_day_index, 10006);

  CB_ASSERT_OK_AND_ASSIGN(std::vector<TimeInfo> backfill,
                          manager.CalculateBackfill(TimeInfo::FromDayIndex(kStartDayIndex),
                                                    end_time, converter, metric, true));
  // Only the previous day index, 10005, is permitted by the 2-day backfill window.
  EXPECT_EQ(ToDayIndices(backfill), std::vector<uint32_t>({10005}));
}

}  // namespace cobalt::local_aggregation