#pragma once

#include <sys/eventfd.h>

#include <array>
#include <atomic>
#include <cassert>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <optional>
#include <queue>
#include <span>
#include <string>
#include <string_view>
#include <thread>
#include <tuple>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <variant>
#include <vector>

#include "absl/base/attributes.h"
#include "absl/base/thread_annotations.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/functional/any_invocable.h"
#include "absl/functional/bind_front.h"
#include "absl/log/log.h"
#include "absl/meta/type_traits.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "boost/asio/error.hpp"
#include "boost/asio/io_context.hpp"
#include "boost/asio/posix/basic_stream_descriptor.hpp"
#include "boost/asio/steady_timer.hpp"
#include "boost/system/error_code.hpp"
#include "subscription.h"
#include "nlohmann/json.hpp"
#include "dbus_utility.hpp" // NOLINT
#include "async_resp.hpp" // NOLINT
#include "http_request.hpp" // NOLINT
#include "managed_store_types.hpp"
#include "sdbusplus/asio/connection.hpp"
#include "sdbusplus/bus/match.hpp"
#include "sdbusplus/message/native_types.hpp"
#include "sdbusplus/message/types.hpp"

namespace managedStore {

// This is a helper class to ensure that a callable is only called once.
// This is useful for constructing a std::function from absl::AnyInvocable.
//
// Example:
//
// absl::AnyInvocable<void(int)> f = ...;
// std::function<void(int)> g = CallableOnce(std::move(f));

// Reference: http://google3/util/functional/callable_once.h
// We cannot use the google3 version because it is not open-sourced.

namespace internal {

// Encapsulates a functor and whether it was called or not.
template <class Functor>
struct CalledState {
  template <class... Args>
  explicit CalledState([[maybe_unused]] int dummy, Args&&... args)  // NOLINT
      : functor(absl::bind_front(std::forward<Args>(args)...)) {}

  std::atomic<bool> called{false};
  Functor functor;
};

// Encapsulates a functor and whether it was called or not, it also check fails
// at destruction if never called.
template <class Functor>
struct CheckCalledState : public CalledState<Functor> {
  using CalledState<Functor>::CalledState;
  ~CheckCalledState() {
    [[maybe_unused]] const bool check =
        this->called.load(std::memory_order_relaxed);
    assert(check && "Functor was never called.");
  }
};

// Encapsulates a Functor state giving it shared semantics and forwards
// operator() const to operator() &&.
template <class Functor, class State>
class SharedCallWrapperAtMostOnce {
 public:
  // The int is to limit overtriggering and to ensure that the copy/move
  // constructors are called when we want to copy/move, rather than this one.
  template <class... Args>
  explicit SharedCallWrapperAtMostOnce([[maybe_unused]] int dummy,
                                       Args&&... args)  // NOLINT
      : internal_(std::make_shared<State>(0, std::forward<Args>(args)...)) {}

  // We use the absl::void_t and decltype(auto) to avoid having the
  // decltype() expression of the return type as part of the mangled name.
  template <class... Args,
            typename = ::absl::void_t<
                decltype(std::declval<Functor>()(std::declval<Args>()...))>>
  decltype(auto) operator()(Args&&... args) const {  // NOLINT
    [[maybe_unused]] const bool called =
        internal_->called.exchange(true, std::memory_order_relaxed);
    assert(!called && "Functor was already called");
    return std::move(internal_->functor)(std::forward<Args>(args)...);
  }

 private:
  std::shared_ptr<State> internal_;
};

template <class... Args>
using CallAtMostOnceT = internal::SharedCallWrapperAtMostOnce<
    decltype(absl::bind_front(std::declval<Args>()...)),
    internal::CalledState<decltype(absl::bind_front(std::declval<Args>()...))>>;

template <typename Functor>
CallAtMostOnceT<Functor> CallAtMostOnce(Functor&& functor) {
  return CallAtMostOnceT<Functor>(0, std::forward<Functor>(functor));
}

template <class... Args>
using CallExactlyOnceT = internal::SharedCallWrapperAtMostOnce<
    decltype(absl::bind_front(std::declval<Args>()...)),
    internal::CheckCalledState<decltype(absl::bind_front(
        std::declval<Args>()...))>>;

template <class... Args>
CallExactlyOnceT<Args...> CallExactlyOnce(Args&&... args) {  // NOLINT
  return CallExactlyOnceT<Args...>(0, std::forward<Args>(args)...);
}

}  // namespace internal

constexpr std::string_view serializationFilePath =
    "/tmp/gBMCwebManagedStore.json";

// forward declaration:
class ManagedObjectStoreHttp;

// ManagedObjectStore caches all responses to GetManagedObjects in memory.
// This store aims to provide O(1) access to ManagedObjectType objects to save
// the round trip latency to query data from each DBus connection.
// A background async task using boost::asio::steady_timer ensures freshness of
// managed objects in store.
class ManagedObjectStore {
 public:
  friend class ManagedObjectStoreHttp;

  explicit ManagedObjectStore(
      const ManagedObjectStoreConfig& cfg, boost::asio::io_context& io,
      std::thread::id main_thread_id,
      const std::shared_ptr<boost::asio::io_context>& io_context_worker_threads,
      ecclesia::SubscriptionService* ss = nullptr,
      sdbusplus::asio::connection* systemBus = nullptr)
      : config(cfg),
        timer(io),
        managedStoreTracker(cfg.pendingDbusResponsesMax),
        time_trace_array_(nlohmann::json::array()),
        subscription_service_(ss),
        io_context_main_thread_(io),
        main_thread_id_(main_thread_id),
        io_context_worker_threads_(io_context_worker_threads),
        system_bus_(systemBus) {}

  virtual ~ManagedObjectStore() {}

  using GetManagedObjectsCb =
      absl::AnyInvocable<void(const boost::system::error_code&,
                              const dbus::utility::ManagedObjectType&)>;

  using GetManagedMapperObjectCb = absl::AnyInvocable<void(
      const boost::system::error_code&, const dbus::utility::MapperGetObject&)>;

  using GetManagedSystemdListUnitsCb =
      absl::AnyInvocable<void(const boost::system::error_code&,
                              const dbus::utility::SystemdListUnits&)>;

  using GetManagedPropertyMapCb =
      absl::AnyInvocable<void(const boost::system::error_code&,
                              const ::dbus::utility::DBusPropertiesMap&)>;

  using GetManagedPropertyCb =
      absl::AnyInvocable<void(const boost::system::error_code&,
                              const dbus::utility::DbusVariantType&, uint64_t)>;

  using GetManagedSubtreeCb =
      absl::AnyInvocable<void(const boost::system::error_code&,
                              const dbus::utility::MapperGetSubTreeResponse&)>;

  using GetManagedSubtreePathsCb = absl::AnyInvocable<void(
      const boost::system::error_code&,
      const dbus::utility::MapperGetSubTreePathsResponse&)>;

  using SetPropertyCb =
      absl::AnyInvocable<void(const boost::system::error_code&)>;

  // Callback used to move managed objects/properties within the store.
  // This allows the store to operate on ValueType instead of specific
  // dbus::utility::DbusVariantType, dbus::utility::ManagedObjectType.
  using ManagedStoreCb =
      absl::AnyInvocable<void(const std::shared_ptr<ValueType>&)>;

  // Maps ManagedStore key to request callback.
  using ManagedStoreCallbackMap =
      std::unordered_map<std::string, std::queue<ManagedStoreCb>>;

  // call the 'callback' with cached/fetched object
  void getManagedObjectsWithContext(
      const std::string& service, const sdbusplus::message::object_path& path,
      const ManagedObjectStoreContext& requestContext,
      GetManagedObjectsCb callback);

  // call the 'callback' with cached/fetched property
  void getProperty(const KeyType& keyType,
                   const ManagedObjectStoreContext& requestContext,
                   GetManagedPropertyCb callback);

  void getDbusObject(const std::string& path,
                     std::span<const std::string_view> interfaces,
                     const ManagedObjectStoreContext& requestContext,
                     GetManagedMapperObjectCb callback);

  void listUnits(const ManagedObjectStoreContext& requestContext,
                 GetManagedSystemdListUnitsCb callback);

  void getAllProperties(const std::string& service,
                        const std::string& objectPath,
                        const std::string& interface,
                        const ManagedObjectStoreContext& requestContext,
                        GetManagedPropertyMapCb callback);

  void getSubTree(const std::string& path, int32_t depth,
                  std::span<const std::string_view> interfaces,
                  const ManagedObjectStoreContext& requestContext,
                  GetManagedSubtreeCb callback);

  void getSubTree(const std::string& path, int32_t depth,
                  const std::vector<std::string>& interfaces,
                  const ManagedObjectStoreContext& requestContext,
                  GetManagedSubtreeCb callback);

  void getSubTreePaths(const std::string& path, int32_t depth,
                       std::span<const std::string_view> interfaces,
                       const ManagedObjectStoreContext& requestContext,
                       GetManagedSubtreePathsCb callback);

  void getAssociatedSubTree(
      const sdbusplus::message::object_path& associatedPath,
      const sdbusplus::message::object_path& path, int32_t depth,
      const std::vector<std::string>& interfaces,
      const ManagedObjectStoreContext& requestContext,
      GetManagedSubtreeCb callback);

  void getAssociatedSubTree(
      const sdbusplus::message::object_path& associatedPath,
      const sdbusplus::message::object_path& path, int32_t depth,
      std::span<const std::string_view> interfaces,
      const ManagedObjectStoreContext& requestContext,
      GetManagedSubtreeCb callback);

  void getAssociatedSubTreePaths(
      const sdbusplus::message::object_path& associatedPath,
      const sdbusplus::message::object_path& path, int32_t depth,
      std::span<const std::string_view> interfaces,
      const ManagedObjectStoreContext& requestContext,
      GetManagedSubtreePathsCb callback);

  template <typename PropertyType>
  void setProperty(const std::string& service, const std::string& path,
                   const std::string& interface,
                   const std::string& propertyName,
                   PropertyType&& propertyValue, SetPropertyCb&& callback) {
    boost::asio::post(
        io_context_main_thread_,
        [this, service, path, interface, propertyName,
         propertyValue(std::forward<PropertyType>(propertyValue)),
         callback(std::forward<SetPropertyCb>(callback))]() mutable {
          system_bus_->async_method_call(
              std::forward<decltype(callback)>(callback), service, path,
              "org.freedesktop.DBus.Properties", "Set", interface, propertyName,
              std::variant<std::decay_t<decltype(propertyValue)>>(
                  std::forward<decltype(propertyValue)>(propertyValue)));
        });
  }

  template <typename Callback, typename... DbusCallArgs>
  void PostDbusCallToIoContextThreadSafe(
      [[maybe_unused]] const std::shared_ptr<boost::asio::io_context::strand>&
          strand,
      Callback&& callback, const std::string& service,
      const std::string& objpath, const std::string& interf,
      const std::string& method, DbusCallArgs&&... dbusCallArgs) {
    // Since everything is in a strand, we don't have to worry about
    // data races.
    boost::asio::post(
        io_context_main_thread_,
        [this, callback(std::forward<Callback>(callback)), service, objpath,
         interf, method,
         ... dbusCallArgs(std::forward<DbusCallArgs>(dbusCallArgs))]() mutable {
          system_bus_->async_method_call(
              std::forward<decltype(callback)>(callback), service, objpath,
              interf, method,
              std::forward<decltype(dbusCallArgs)>(dbusCallArgs)...);
        });
  }

  boost::asio::io_context& GetIoContext() const {
    return io_context_main_thread_;
  }

  boost::asio::io_context& GetWorkerThreadsIoContext() const {
    return *io_context_worker_threads_;
  }

  bool IsInMainThread() const {
    return std::this_thread::get_id() == main_thread_id_;
  }

  // DO NOT USE THIS METHOD IF YOU ARE HANDLING A GET REQUEST
  // Get requests will be multithreaded and using system bus in multiple threads
  // will lead to a data race.
  // This method only exists for legacy code.
  // NOTE: currently only get requests are multithreaded
  // system bus usage should be minimized
  sdbusplus::asio::connection* GetDeprecatedThreadUnsafeSystemBus() const {
    return system_bus_;
  }

  ManagedStoreTracker getStoreTracker() const {
    absl::MutexLock lock(&managed_store_tracker_mutex_);
    return this->managedStoreTracker;
  }

  // TODO(rahulkpr): migrate calls to subscribeToInterfacesAdded().
  void subscribe([[maybe_unused]] const ManagedObjectStoreContext& context,
                 [[maybe_unused]] const std::string& path,
                 [[maybe_unused]] const std::string& interface = "") {}

  // Subscribe to `InterfacesAdded` signal from D-Bus.
  void subscribeToInterfacesAdded(
      const std::string& object_path,
      const ManagedObjectStoreContext& request_context);

  // check if the store is enabled:
  bool isEnabled() const { return config.isEnabled; }

  // get the current config
  const ManagedObjectStoreConfig& getConfig() const { return config; }

  void storeTimeTrace(const nlohmann::json& tTrace);

  void clearTimeTrace();

  bool serialize(std::string_view filePath = serializationFilePath);

  bool deserialize(std::string_view filePath = serializationFilePath);

  bool toJson(nlohmann::json& obj);

  bool fromJson(const nlohmann::json& obj);

  nlohmann::json GetSubscriptionsToJSON();

  nlohmann::json GetEventsToJSON();

  nlohmann::json GetDBusMonitorsToJSON();

  nlohmann::json GetEventsBySubscriptionIdToJSON(size_t subscription_id);

  nlohmann::json ClearEventStore();

  // protected APIs (SPI):
 protected:
  // We want to unify PostDbusCallToIoContextThreadSafe into one method such
  // that we dont expose a threadsafe and non-threadsafe method. However, we
  // need to be able to distinguish between different dbus calls. Get calls can
  // be multithreaded, but all other calls will not be multithreaded. This means
  // we need a mechanism at compile time to distinguish between different dbus
  // calls. We can use the return type of the dbus call to do this.

  // The dbus calls that can be multithreaded are:
  // 1. The dbus callback must have 2 parameters
  // 2. sdbusplus can compile the return type

  // If these conditions are not met, we cant not use template deduction for
  // callback return type or else compilation will fail.

  template <typename Callback, typename... DbusCallArgs>
  static constexpr bool DbusCallbackIsCompilable() {
    // All Dbus Get calls have 2 args:
    // 1. error_code
    // 2. return value
    if constexpr (std::tuple_size<
                      boost::callable_traits::args_t<Callback>>::value != 2) {
      return false;
    } else {
      using CallbackReturnType = std::decay_t<
          std::tuple_element_t<1, boost::callable_traits::args_t<Callback>>>;

      using undefined_type_id =
          sdbusplus::message::types::details::undefined_type_id;

      // Sdbus requires dbus object return types to be defined.
      // This static assert is done here:
      // https://github.com/openbmc/sdbusplus/blob/master/include/sdbusplus/message/types.hpp#L278
      // We need to do it here as the yield_method_call requires the return type
      // to be defined.
      return !std::is_base_of_v<
          undefined_type_id,
          sdbusplus::message::types::details::type_id<CallbackReturnType>>;
    }
  }

  // clear/flush all managed objects:
  void clearAllObjects();

  // Subscribe to DBus signals with specific `match_expression`.
  // Uses given `event_source_id` to associate D-Bus match created.
  void subscribe(const ManagedObjectStoreContext& requestContext,
                 const std::string& match_expression,
                 ecclesia::EventSourceId event_source_id);

  // Calls into dbus service to get managed objects.
  // Virtual to allow unit testing and possible derived implementations of
  // managedStore.
  virtual void getManagedObjectsFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      const std::string& service, const sdbusplus::message::object_path& path,
      ManagedStoreCb callback, std::optional<uint64_t> refresh_interval);

  // Calls into dbus service to get Dbus object.
  virtual void getManagedMapperObjectFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      const sdbusplus::message::object_path& path,
      const std::vector<std::string>& interfaces, ManagedStoreCb callback,
      std::optional<uint64_t> refresh_interval);

  // Calls into dbus service to get managed properties.
  virtual void getManagedPropertiesFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      const std::string& service, const sdbusplus::message::object_path& path,
      const std::string& interface, const std::string& property,
      ManagedStoreCb callback, std::optional<uint64_t> refresh_interval);

  // Calls into dbus service to getAllProperties.
  virtual void getManagedPropertiesMapFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      const std::string& service, const sdbusplus::message::object_path& path,
      const std::string& interface, ManagedStoreCb callback,
      std::optional<uint64_t> refresh_interval);

  // Calls into dbus service to get subtree.
  virtual void getManagedSubtreeFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      const sdbusplus::message::object_path& path, int32_t depth,
      const std::vector<std::string>& interfaces, ManagedStoreCb callback,
      std::optional<uint64_t> refresh_interval);

  // Calls into dbus service to get subtree paths.
  virtual void getManagedSubtreePathsFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      const sdbusplus::message::object_path& path, int32_t depth,
      const std::vector<std::string>& interfaces, ManagedStoreCb callback,
      std::optional<uint64_t> refresh_interval);

  virtual void getManagedAssociatedSubtreeFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      const sdbusplus::message::object_path& associatedPath,
      const sdbusplus::message::object_path& path, int32_t depth,
      const std::vector<std::string>& interfaces, ManagedStoreCb callback,
      std::optional<uint64_t> refresh_interval);

  virtual void getManagedAssociatedSubtreePathsFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      const sdbusplus::message::object_path& associatedPath,
      const sdbusplus::message::object_path& path, int32_t depth,
      const std::vector<std::string>& interfaces, ManagedStoreCb callback,
      std::optional<uint64_t> refresh_interval);

  virtual void getSystemdUnitsFromDbusService(
      const std::shared_ptr<boost::asio::io_context::strand>& strand,
      ManagedStoreCb callback, std::optional<uint64_t> refresh_interval);

  // These are in protected because the subclass needs to access them for unit
  // tests.
  // Stores ManagedObjectType objects with D-Bus interface and connection
  // as key along with other metadata necessary for scheduling data refresh.
  mutable absl::Mutex managed_objects_mutex_;
  ManagedObjectsMap managedObjects ABSL_GUARDED_BY(managed_objects_mutex_);

  // schedule the next refresh event (timer):
  void scheduleRefresh();

  // fetch the given object from DBUS:
  // Refresh operation is considered high priority when isReadThroughRequest
  // is set to true i.e the refresh is not subject to the pending i/o
  // threshold invariant.
  // If `strand` is set, the callback will be posted to its original thread.
  void refreshObject(
      const KeyType& key, bool isReadThroughRequest = false,
      const std::shared_ptr<boost::asio::io_context::strand>& strand = nullptr);

  // Process dbus response after refreshObject(), invokes given callback if
  // valid.
  void processDbusResponse(const KeyType& keyType,
                           const std::shared_ptr<ValueType>& managedValue,
                           bool isReadThroughRequest,
                           std::optional<uint64_t> refresh_interval);

  // Get managed entities (properties / objects) from store.
  // Triggers refresh if not in store.
  void getStoredOrReadThrough(const KeyType& keyType,
                              const ManagedObjectStoreContext& requestContext,
                              ManagedStoreCb callback);

  // Compares two managed store values and returns true if they are different
  static bool DetectChange(const ValueType& valueA, const ValueType& valueB);

  // Stores keys that have custom refresh interval.
  absl::Mutex keys_to_custom_interval_mutex_;
  absl::flat_hash_map<std::string, uint64_t> keys_to_custom_interval_
      ABSL_GUARDED_BY(keys_to_custom_interval_mutex_);

  // `pq` is only accessed by the io_context thread.
  // priority queue of refreshes:
  ManagedStorePriorityQueue pq;

  // Tracks callback for incoming requests.
  absl::Mutex request_cb_map_mutex_;
  ManagedStoreCallbackMap requestCbMap ABSL_GUARDED_BY(request_cb_map_mutex_);

  // `key_to_pq_entry_` is only accessed by the io_context thread.
  // Tracks number of priority queue entries per managedStore key.
  absl::flat_hash_map<std::string, size_t> key_to_pq_entry_;

  // runtime config for the store
  const ManagedObjectStoreConfig config;

  // `timer` is only accessed by the io_context thread.
  // A single timer object is used to asynchronously wait for a fixed duration
  // before refreshing ManagedObject in store with the shortest deadline.
  boost::asio::steady_timer timer;

  // Flag set to true when scheduler has scheduled a refresh operation.
  bool isSchedulerActive = false;

  // Tracks state of managed object store.
  mutable absl::Mutex managed_store_tracker_mutex_
      ABSL_ACQUIRED_AFTER(managed_objects_mutex_);
  ManagedStoreTracker managedStoreTracker
      ABSL_GUARDED_BY(managed_store_tracker_mutex_);

  // Mutex to protect signal_monitors_
  absl::Mutex signal_monitor_mutex_;

  // Tracks signal monitors
  absl::flat_hash_map<ecclesia::EventSourceId,
                      std::unique_ptr<sdbusplus::bus::match_t>>
      signal_monitors_ ABSL_GUARDED_BY(signal_monitor_mutex_);

  // Time Trace
  mutable absl::Mutex time_trace_array_mutex_;
  nlohmann::json time_trace_array_ ABSL_GUARDED_BY(time_trace_array_mutex_);

  // Pointer to a SubscriptionService that's used to publish a change in
  // an observed resource via the service's Notify callback.
  // Not owned.
  ecclesia::SubscriptionService* subscription_service_;

  // The main thread that's runs dbus calls, priority queue refresh.
  boost::asio::io_context& io_context_main_thread_;
  // The thread id of the main thread
  const std::thread::id main_thread_id_;
  // The worker threads that's runs processing of cache hits.
  // Only used when `multi_thread_get` is set to true.
  std::shared_ptr<boost::asio::io_context> io_context_worker_threads_;

  sdbusplus::asio::connection* system_bus_;

  // Deprecated: dangerous and not thread-safe should not be used other than
  // unit tests.
  ABSL_DEPRECATED("Use GetManagedObjectsMetrics instead")
  const ManagedObjectsMap& getManagedObjects() const
      ABSL_LOCKS_EXCLUDED(managed_objects_mutex_) {
    absl::MutexLock lock(&managed_objects_mutex_);
    return managedObjects;
  }

  // Returns a json object with metrics about the managed objects.
  nlohmann::json GetManagedObjectsMetrics(
      std::chrono::steady_clock::time_point now) const
      ABSL_LOCKS_EXCLUDED(managed_objects_mutex_);

  // Returns the estimated size in bytes of the cached states.
  std::size_t GetSnapshotSizeInBytes() const
      ABSL_LOCKS_EXCLUDED(managed_objects_mutex_);

  // Updates an object with updated time points and refresh status if it already
  // exists within managed store. Otherwise, it adds a new object to managed
  // store.
  void UpdateOrInsertManagedObjects(const std::string& key_id,
                                    const std::shared_ptr<ValueType>& value);

  // Increments numRefreshScheduled value of an object with thread-safety if it
  // exists.
  void UpdateNumRefreshesScheduled(std::string& key_id);

  // Checks if an object is still tracked within managed store.
  bool IsManagedObjectTracked(const std::string& key_id) {
    absl::MutexLock lock(&managed_objects_mutex_);
    return managedObjects.find(key_id) != managedObjects.end();
  }

  // Evicts an object by key_id from managed store if it exists with
  // thread-safety.
  void EvictManagedObject(const std::string& key_id) {
    absl::MutexLock lock(&managed_objects_mutex_);
    // Early return if Managed Object is no longer tracked.
    if (managedObjects.find(key_id) == managedObjects.end()) return;
    // Only evict Managed Object if it is still tracked.
    managedObjects.erase(key_id);
  }

  // Checks if an object is current and still tracked within managed store with
  // thread-safety.
  // If true, the cached object shared pointer is copied into `copied_object`,
  // which extends the lifecycle. If false, the stabled object is erased from
  // the store, and nullptr is returned.
  virtual std::shared_ptr<ValueType> IsManagedObjectCurrentAndTracked(
      const std::string& key_id, const KeyType& keyType,
      const ManagedObjectStoreContext& requestContextValueType);

  std::optional<uint64_t> GetRefreshInterval(const std::string& key_id)
      ABSL_LOCKS_EXCLUDED(keys_to_custom_interval_mutex_) {
    // Note: we might get a outdated refresh interval here, but it's ok. The
    // next call will use the updated one.
    absl::MutexLock lock(&keys_to_custom_interval_mutex_);
    auto it = keys_to_custom_interval_.find(key_id);
    if (it == keys_to_custom_interval_.end()) {
      return std::nullopt;
    }
    return it->second;
  }

  void SetRefreshInterval(const std::string& key_id, uint64_t refresh_interval)
      ABSL_LOCKS_EXCLUDED(keys_to_custom_interval_mutex_) {
    absl::MutexLock lock(&keys_to_custom_interval_mutex_);
    keys_to_custom_interval_[key_id] = refresh_interval;
  }

  // Returns false if the `key_id` is already in the map which means no refresh
  // call is required.
  // The combination of `InsertCallbackIntoRequestCbMapAndReturnTrueIfNew` and
  // `GetCallbacksFromRequestCbMap` should make sure no callbacks are lost.
  bool InsertCallbackIntoRequestCbMapAndReturnTrueIfNew(
      const std::string& key_id, ManagedStoreCb callback)
      ABSL_LOCKS_EXCLUDED(request_cb_map_mutex_) {
    absl::MutexLock lock(&request_cb_map_mutex_);
    auto cb_map_iter = requestCbMap.find(key_id);
    if (cb_map_iter != requestCbMap.end()) {
      // Request collapsing - Combine multiple requests for the same dbus object
      // into a single request and use the resulting response to satisfy all
      // duplicate requests.
      cb_map_iter->second.push(std::move(callback));
      return false;
    }
    requestCbMap[key_id].push(std::move(callback));
    return true;
  }

  std::queue<ManagedStoreCb> GetCallbacksFromRequestCbMap(
      const std::string& key_id) ABSL_LOCKS_EXCLUDED(request_cb_map_mutex_) {
    std::queue<ManagedStoreCb> callbacks;
    absl::MutexLock lock(&request_cb_map_mutex_);
    auto cb_map_iter = requestCbMap.find(key_id);
    if (cb_map_iter != requestCbMap.end()) {
      callbacks = std::move(cb_map_iter->second);
      requestCbMap.erase(cb_map_iter);
    }
    return callbacks;
  }

  // Deprecated: dangerous and not thread-safe should not be used other than
  // unit tests.
  ABSL_DEPRECATED("Use GetPriorityQueueMetrics instead")
  const ManagedStorePriorityQueue& getPriorityQueue() const { return pq; }

  // Fetches the current priority queue, generates a json object and feeds it
  // into the callback.
  void GetPriorityQueueMetrics(
      absl::AnyInvocable<void(const nlohmann::json&)> callback) {
    // Priority queue should only be accessed by the io_context thread.
    boost::asio::post(
        io_context_main_thread_,
        [this, callback = std::move(callback)]() mutable {
          auto jsonObjPqArray = nlohmann::json::array();
          size_t entries = 0;
          // Max number of PQ Entries that should be exported for debug.
          constexpr size_t kPQExportLimit = 20000;
          ManagedStorePriorityQueue pq_copy = pq;
          while (!pq_copy.empty() && entries++ < kPQExportLimit) {
            const auto& entry = pq_copy.top();
            jsonObjPqArray.push_back(entry.managedStoreKey.toString());
            pq_copy.pop();
          }
          std::move(callback)(jsonObjPqArray);
        });
  }

  // Inserts a new entry into the priority queue. Should only be called by the
  // io_context thread and used in unit tests.
  void InsertPriorityQueueEntry(const PriorityQueueEntry& entry) {
    // Priority queue should only be accessed by the io_context thread.
    boost::asio::post(io_context_main_thread_,
                      [this, &entry]() { pq.push(entry); });
  }

  void GetPriorityQueueSize(absl::AnyInvocable<void(size_t)> callback) {
    // Priority queue should only be accessed by the io_context thread.
    boost::asio::post(io_context_main_thread_,
                      [this, callback = std::move(callback)]() mutable {
                        std::move(callback)(pq.size());
                      });
  }

  nlohmann::json getTimeTrace() const
      ABSL_LOCKS_EXCLUDED(time_trace_array_mutex_) {
    absl::MutexLock lock(&time_trace_array_mutex_);
    return time_trace_array_;
  }

  // Disable copy constructors (singelton)
  ManagedObjectStore(const ManagedObjectStore&) = delete;
  ManagedObjectStore& operator=(const ManagedObjectStore&) = delete;

  // Not yet: https://gbmc-private-review.git.corp.google.com/c/gbmcweb/+/3520
  // ManagedObjectStore(ManagedObjectStore&&) = delete;
  // ManagedObjectStore& operator=(ManagedObjectStore&&) = delete;

  // TODO: use this macro / template ^
  // https://www.boost.org/doc/libs/1_81_0/libs/serialization/doc/singleton.html
};

#ifdef UNIT_TEST_BUILD
// Forward Declaration
class MockSerializedManagedObjectStore;
using ManagedStore = managedStore::MockSerializedManagedObjectStore;
#else
using ManagedStore = managedStore::ManagedObjectStore;
#endif

// Only initialize if cfg and io are not null
// TODO(rahulkpr): Call InitializeManagedStore() instead from webserver_main.
ManagedStore* InitializeManagedStore(
    const ManagedObjectStoreConfig* cfg,
    boost::asio::io_context* io_main_thread,
    const std::shared_ptr<boost::asio::io_context>& io_context_worker_threads,
    ecclesia::SubscriptionService* subscription_service,
    sdbusplus::asio::connection* systemBus);

// If ManagedObjectStore is initialized, then return it, if not return nullptr
ManagedStore* GetManagedObjectStore();

}  // namespace managedStore