// Copyright 2024 The Pigweed Authors
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not
// use this file except in compliance with the License. You may obtain a copy of
// the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
// License for the specific language governing permissions and limitations under
// the License.

#include "modules/pubsub/pubsub.h"

#include "pw_async2/dispatcher.h"
#include "pw_unit_test/framework.h"

namespace {

using namespace std::literals::chrono_literals;

// Test fixtures.

struct EchoRequest {
  uint32_t value;
};

struct EchoResponse {
  void AddValue(uint32_t v) {
    value += v;
    ++events_seen;
  }

  uint32_t value = 0;
  size_t events_seen = 0;
};

class PubSubTest : public ::testing::Test {
 protected:
  static constexpr size_t kMaxEvents = 4;
  static constexpr size_t kMaxSubscribers = 4;
  using PubSub =
      sense::GenericPubSubBuffer<EchoRequest, kMaxEvents, kMaxSubscribers>;

  void SetUp() override { pubsub_.Init(dispatcher_); }

  pw::async2::Dispatcher dispatcher_;
  PubSub pubsub_;
  std::array<EchoResponse, kMaxSubscribers> responses_;
};

// Unit tests.

TEST_F(PubSubTest, Publish_OneSubscriber) {
  EchoResponse& response = responses_[0];
  ASSERT_TRUE(pubsub_.Subscribe(
      [&response](EchoRequest request) { response.AddValue(request.value); }));
  EXPECT_TRUE(pubsub_.Publish({.value = 42u}));
  EXPECT_EQ(dispatcher_.RunUntilStalled(), pw::async2::Pending());
  EXPECT_EQ(response.value, 42u);
}

TEST_F(PubSubTest, Publish_MultipleSubscribers) {
  for (auto& response : responses_) {
    ASSERT_TRUE(pubsub_.Subscribe([&response](EchoRequest request) {
      response.AddValue(request.value);
    }));
  }
  ASSERT_TRUE(pubsub_.Publish({.value = 4u}));
  EXPECT_EQ(dispatcher_.RunUntilStalled(), pw::async2::Pending());
  for (auto& response : responses_) {
    EXPECT_EQ(response.value, 4u);
  }
}

TEST_F(PubSubTest, Publish_MultipleEvents) {
  EchoResponse& response = responses_[0];
  ASSERT_TRUE(pubsub_.Subscribe(
      [&response](EchoRequest request) { response.AddValue(request.value); }));
  ASSERT_TRUE(pubsub_.Publish({.value = 1}));
  ASSERT_TRUE(pubsub_.Publish({.value = 2}));
  ASSERT_TRUE(pubsub_.Publish({.value = 3}));
  ASSERT_TRUE(pubsub_.Publish({.value = 4}));
  EXPECT_EQ(dispatcher_.RunUntilStalled(), pw::async2::Pending());
  EXPECT_EQ(response.value, 10u);

  ASSERT_TRUE(pubsub_.Publish({.value = 5}));
  ASSERT_TRUE(pubsub_.Publish({.value = 6}));
  ASSERT_TRUE(pubsub_.Publish({.value = 7}));
  ASSERT_TRUE(pubsub_.Publish({.value = 8}));
  EXPECT_EQ(dispatcher_.RunUntilStalled(), pw::async2::Pending());
  EXPECT_EQ(response.value, 36u);
}

TEST_F(PubSubTest, Publish_MultipleEvents_QueueFull) {
  EchoResponse& response = responses_[0];

  ASSERT_TRUE(pubsub_.Subscribe(
      [&response](EchoRequest request) { response.AddValue(request.value); }));
  // Note: event queue only has room for 4 events
  ASSERT_TRUE(pubsub_.Publish({.value = 10}));
  ASSERT_TRUE(pubsub_.Publish({.value = 11}));
  ASSERT_TRUE(pubsub_.Publish({.value = 12}));
  ASSERT_TRUE(pubsub_.Publish({.value = 13}));
  EXPECT_FALSE(pubsub_.Publish({.value = 14}));

  // The fifth event never gets sent.
  EXPECT_EQ(dispatcher_.RunUntilStalled(), pw::async2::Pending());
  EXPECT_EQ(response.events_seen, 4u);
  EXPECT_EQ(response.value, 46u);
}

TEST_F(PubSubTest, Subscribe_Full) {
  for (auto& response : responses_) {
    ASSERT_TRUE(pubsub_.Subscribe([&response](EchoRequest request) {
      response.AddValue(request.value);
    }));
  }

  // Subscriber is not added when max subscribers has been reached.
  EchoResponse extra;
  EXPECT_FALSE(pubsub_.Subscribe(
      [&extra](EchoRequest request) { extra.AddValue(request.value); }));
  EXPECT_EQ(pubsub_.subscriber_count(), 4u);

  // The extra subscriber never gets events.
  ASSERT_TRUE(pubsub_.Publish({.value = 4u}));
  EXPECT_EQ(dispatcher_.RunUntilStalled(), pw::async2::Pending());
  EXPECT_EQ(extra.events_seen, 0u);
}

TEST_F(PubSubTest, Subscribe_Unsubscribe) {
  std::array<PubSub::SubscribeToken, kMaxSubscribers> tokens;
  auto token = tokens.begin();
  for (auto& response : responses_) {
    ASSERT_NE(token, tokens.end());
    auto result = pubsub_.Subscribe(
        [&response](EchoRequest request) { response.AddValue(request.value); });
    EXPECT_TRUE(result.has_value());
    *token = *result;
    ++token;
  }
  EXPECT_EQ(pubsub_.subscriber_count(), 4u);

  // Remove a subscriber.
  EXPECT_TRUE(pubsub_.Unsubscribe(tokens[1]));
  EXPECT_EQ(pubsub_.subscriber_count(), 3u);

  // Add a subscriber after removing.
  EchoResponse& response = responses_[2];
  EXPECT_TRUE(pubsub_.Subscribe(
      [&response](EchoRequest request) { response.AddValue(request.value); }));
  EXPECT_EQ(pubsub_.subscriber_count(), 4u);

  // Remove multiple subscriber.
  EXPECT_TRUE(pubsub_.Unsubscribe(tokens[0]));
  EXPECT_TRUE(pubsub_.Unsubscribe(tokens[2]));
  EXPECT_TRUE(pubsub_.Unsubscribe(tokens[3]));
  EXPECT_EQ(pubsub_.subscriber_count(), 1u);

  // Unsubscribe invalid.
  EXPECT_FALSE(pubsub_.Unsubscribe(tokens[1]));
}

}  // namespace