// Copyright 2023 The Chromium 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 "quiche/common/btree_scheduler.h"

#include <optional>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/types/span.h"
#include "quiche/common/platform/api/quiche_test.h"
#include "quiche/common/test_tools/quiche_test_utils.h"

namespace quiche::test {
namespace {

using ::testing::ElementsAre;
using ::testing::Optional;

template <typename Id, typename Priority>
void ScheduleIds(BTreeScheduler<Id, Priority>& scheduler,
                 absl::Span<const Id> ids) {
  for (Id id : ids) {
    QUICHE_EXPECT_OK(scheduler.Schedule(id));
  }
}

template <typename Id, typename Priority>
std::vector<Id> PopAll(BTreeScheduler<Id, Priority>& scheduler) {
  std::vector<Id> result;
  result.reserve(scheduler.NumScheduled());
  for (;;) {
    absl::StatusOr<Id> id = scheduler.PopFront();
    if (id.ok()) {
      result.push_back(*id);
    } else {
      EXPECT_THAT(id, StatusIs(absl::StatusCode::kNotFound));
      break;
    }
  }
  return result;
}

TEST(BTreeSchedulerTest, SimplePop) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(1, 100));
  QUICHE_EXPECT_OK(scheduler.Register(2, 101));
  QUICHE_EXPECT_OK(scheduler.Register(3, 102));

  EXPECT_THAT(scheduler.GetPriorityFor(1), Optional(100));
  EXPECT_THAT(scheduler.GetPriorityFor(3), Optional(102));
  EXPECT_EQ(scheduler.GetPriorityFor(5), std::nullopt);

  EXPECT_EQ(scheduler.NumScheduled(), 0u);
  EXPECT_FALSE(scheduler.HasScheduled());
  QUICHE_EXPECT_OK(scheduler.Schedule(1));
  QUICHE_EXPECT_OK(scheduler.Schedule(2));
  QUICHE_EXPECT_OK(scheduler.Schedule(3));
  EXPECT_EQ(scheduler.NumScheduled(), 3u);
  EXPECT_TRUE(scheduler.HasScheduled());

  EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(3));
  EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(2));
  EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(1));

  QUICHE_EXPECT_OK(scheduler.Schedule(2));
  QUICHE_EXPECT_OK(scheduler.Schedule(1));
  QUICHE_EXPECT_OK(scheduler.Schedule(3));

  EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(3));
  EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(2));
  EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(1));

  QUICHE_EXPECT_OK(scheduler.Schedule(3));
  QUICHE_EXPECT_OK(scheduler.Schedule(1));

  EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(3));
  EXPECT_THAT(scheduler.PopFront(), IsOkAndHolds(1));
}

TEST(BTreeSchedulerTest, FIFO) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(1, 100));
  QUICHE_EXPECT_OK(scheduler.Register(2, 100));
  QUICHE_EXPECT_OK(scheduler.Register(3, 100));

  ScheduleIds(scheduler, {2, 1, 3});
  EXPECT_THAT(PopAll(scheduler), ElementsAre(2, 1, 3));

  QUICHE_EXPECT_OK(scheduler.Register(4, 101));
  QUICHE_EXPECT_OK(scheduler.Register(5, 99));

  ScheduleIds(scheduler, {5, 1, 2, 3, 4});
  EXPECT_THAT(PopAll(scheduler), ElementsAre(4, 1, 2, 3, 5));
  ScheduleIds(scheduler, {1, 5, 2, 4, 3});
  EXPECT_THAT(PopAll(scheduler), ElementsAre(4, 1, 2, 3, 5));
  ScheduleIds(scheduler, {3, 5, 2, 4, 1});
  EXPECT_THAT(PopAll(scheduler), ElementsAre(4, 3, 2, 1, 5));
  ScheduleIds(scheduler, {3, 2, 1, 2, 3});
  EXPECT_THAT(PopAll(scheduler), ElementsAre(3, 2, 1));
}

TEST(BTreeSchedulerTest, NumEntriesInRange) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(1, 0));
  QUICHE_EXPECT_OK(scheduler.Register(2, 0));
  QUICHE_EXPECT_OK(scheduler.Register(3, 0));
  QUICHE_EXPECT_OK(scheduler.Register(4, -2));
  QUICHE_EXPECT_OK(scheduler.Register(5, -5));
  QUICHE_EXPECT_OK(scheduler.Register(6, 10));
  QUICHE_EXPECT_OK(scheduler.Register(7, 16));
  QUICHE_EXPECT_OK(scheduler.Register(8, 32));
  QUICHE_EXPECT_OK(scheduler.Register(9, 64));

  EXPECT_EQ(scheduler.NumScheduled(), 0u);
  EXPECT_EQ(scheduler.NumScheduledInPriorityRange(std::nullopt, std::nullopt),
            0u);
  EXPECT_EQ(scheduler.NumScheduledInPriorityRange(-1, 1), 0u);

  for (int stream = 1; stream <= 9; ++stream) {
    QUICHE_ASSERT_OK(scheduler.Schedule(stream));
  }

  EXPECT_EQ(scheduler.NumScheduled(), 9u);
  EXPECT_EQ(scheduler.NumScheduledInPriorityRange(std::nullopt, std::nullopt),
            9u);
  EXPECT_EQ(scheduler.NumScheduledInPriorityRange(0, 0), 3u);
  EXPECT_EQ(scheduler.NumScheduledInPriorityRange(std::nullopt, -1), 2u);
  EXPECT_EQ(scheduler.NumScheduledInPriorityRange(1, std::nullopt), 4u);
}

TEST(BTreeSchedulerTest, Registration) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(1, 0));
  QUICHE_EXPECT_OK(scheduler.Register(2, 0));

  QUICHE_EXPECT_OK(scheduler.Schedule(1));
  QUICHE_EXPECT_OK(scheduler.Schedule(2));
  EXPECT_EQ(scheduler.NumScheduled(), 2u);
  EXPECT_TRUE(scheduler.IsScheduled(2));

  EXPECT_THAT(scheduler.Register(2, 0),
              StatusIs(absl::StatusCode::kAlreadyExists));
  QUICHE_EXPECT_OK(scheduler.Unregister(2));
  EXPECT_EQ(scheduler.NumScheduled(), 1u);
  EXPECT_FALSE(scheduler.IsScheduled(2));

  EXPECT_THAT(scheduler.UpdatePriority(2, 1234),
              StatusIs(absl::StatusCode::kNotFound));
  EXPECT_THAT(scheduler.Unregister(2), StatusIs(absl::StatusCode::kNotFound));
  EXPECT_THAT(scheduler.Schedule(2), StatusIs(absl::StatusCode::kNotFound));
  QUICHE_EXPECT_OK(scheduler.Register(2, 0));
  EXPECT_EQ(scheduler.NumScheduled(), 1u);
  EXPECT_TRUE(scheduler.IsScheduled(1));
  EXPECT_FALSE(scheduler.IsScheduled(2));
}

TEST(BTreeSchedulerTest, UpdatePriorityUp) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(1, 0));
  QUICHE_EXPECT_OK(scheduler.Register(2, 0));
  QUICHE_EXPECT_OK(scheduler.Register(3, 0));

  ScheduleIds(scheduler, {1, 2, 3});
  QUICHE_EXPECT_OK(scheduler.UpdatePriority(2, 1000));
  EXPECT_THAT(PopAll(scheduler), ElementsAre(2, 1, 3));
}

TEST(BTreeSchedulerTest, UpdatePriorityDown) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(1, 0));
  QUICHE_EXPECT_OK(scheduler.Register(2, 0));
  QUICHE_EXPECT_OK(scheduler.Register(3, 0));

  ScheduleIds(scheduler, {1, 2, 3});
  QUICHE_EXPECT_OK(scheduler.UpdatePriority(2, -1000));
  EXPECT_THAT(PopAll(scheduler), ElementsAre(1, 3, 2));
}

TEST(BTreeSchedulerTest, UpdatePriorityEqual) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(1, 0));
  QUICHE_EXPECT_OK(scheduler.Register(2, 0));
  QUICHE_EXPECT_OK(scheduler.Register(3, 0));

  ScheduleIds(scheduler, {1, 2, 3});
  QUICHE_EXPECT_OK(scheduler.UpdatePriority(2, 0));
  EXPECT_THAT(PopAll(scheduler), ElementsAre(1, 2, 3));
}

TEST(BTreeSchedulerTest, UpdatePriorityIntoSameBucket) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(1, 0));
  QUICHE_EXPECT_OK(scheduler.Register(2, -100));
  QUICHE_EXPECT_OK(scheduler.Register(3, 0));

  ScheduleIds(scheduler, {1, 2, 3});
  QUICHE_EXPECT_OK(scheduler.UpdatePriority(2, 0));
  EXPECT_THAT(PopAll(scheduler), ElementsAre(1, 2, 3));
}

TEST(BTreeSchedulerTest, ShouldYield) {
  BTreeScheduler<int, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(10, 100));
  QUICHE_EXPECT_OK(scheduler.Register(20, 101));
  QUICHE_EXPECT_OK(scheduler.Register(21, 101));
  QUICHE_EXPECT_OK(scheduler.Register(30, 102));

  EXPECT_THAT(scheduler.ShouldYield(10), IsOkAndHolds(false));
  EXPECT_THAT(scheduler.ShouldYield(20), IsOkAndHolds(false));
  EXPECT_THAT(scheduler.ShouldYield(21), IsOkAndHolds(false));
  EXPECT_THAT(scheduler.ShouldYield(30), IsOkAndHolds(false));
  EXPECT_THAT(scheduler.ShouldYield(40), StatusIs(absl::StatusCode::kNotFound));

  QUICHE_EXPECT_OK(scheduler.Schedule(20));

  EXPECT_THAT(scheduler.ShouldYield(10), IsOkAndHolds(true));
  EXPECT_THAT(scheduler.ShouldYield(20), IsOkAndHolds(false));
  EXPECT_THAT(scheduler.ShouldYield(21), IsOkAndHolds(true));
  EXPECT_THAT(scheduler.ShouldYield(30), IsOkAndHolds(false));
}

struct CustomPriority {
  int a;
  int b;

  bool operator<(const CustomPriority& other) const {
    return std::make_tuple(a, b) < std::make_tuple(other.a, other.b);
  }
};

TEST(BTreeSchedulerTest, CustomPriority) {
  BTreeScheduler<int, CustomPriority> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(10, CustomPriority{0, 1}));
  QUICHE_EXPECT_OK(scheduler.Register(11, CustomPriority{0, 0}));
  QUICHE_EXPECT_OK(scheduler.Register(12, CustomPriority{0, 0}));
  QUICHE_EXPECT_OK(scheduler.Register(13, CustomPriority{10, 0}));
  QUICHE_EXPECT_OK(scheduler.Register(14, CustomPriority{-10, 0}));

  ScheduleIds(scheduler, {10, 11, 12, 13, 14});
  EXPECT_THAT(PopAll(scheduler), ElementsAre(13, 10, 11, 12, 14));
}

struct CustomId {
  int a;
  std::string b;

  bool operator==(const CustomId& other) const {
    return a == other.a && b == other.b;
  }

  template <typename H>
  friend H AbslHashValue(H h, const CustomId& c) {
    return H::combine(std::move(h), c.a, c.b);
  }
};

std::ostream& operator<<(std::ostream& os, const CustomId& id) {
  os << id.a << ":" << id.b;
  return os;
}

TEST(BTreeSchedulerTest, CustomIds) {
  BTreeScheduler<CustomId, int> scheduler;
  QUICHE_EXPECT_OK(scheduler.Register(CustomId{1, "foo"}, 10));
  QUICHE_EXPECT_OK(scheduler.Register(CustomId{1, "bar"}, 12));
  QUICHE_EXPECT_OK(scheduler.Register(CustomId{2, "foo"}, 11));
  EXPECT_THAT(scheduler.Register(CustomId{1, "foo"}, 10),
              StatusIs(absl::StatusCode::kAlreadyExists));

  ScheduleIds(scheduler,
              {CustomId{1, "foo"}, CustomId{1, "bar"}, CustomId{2, "foo"}});
  EXPECT_THAT(scheduler.ShouldYield(CustomId{1, "foo"}), IsOkAndHolds(true));
  EXPECT_THAT(scheduler.ShouldYield(CustomId{1, "bar"}), IsOkAndHolds(false));
  EXPECT_THAT(
      PopAll(scheduler),
      ElementsAre(CustomId{1, "bar"}, CustomId{2, "foo"}, CustomId{1, "foo"}));
}

}  // namespace
}  // namespace quiche::test