#ifndef THIRD_PARTY_MILOTIC_EXTERNAL_CC_HFT_ADAPTER_DATA_SOURCE_H_
#define THIRD_PARTY_MILOTIC_EXTERNAL_CC_HFT_ADAPTER_DATA_SOURCE_H_

#include <cstddef>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/hash/hash.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/substitute.h"
#include "absl/time/time.h"
#include "g3/macros.h"
#include "fru_identifier.pb.h"
#include "identifier.pb.h"
#include "payload.pb.h"
#include "sensor_identifier.pb.h"
#include "subscription_params.pb.h"

namespace milotic_hft {

class SensorMutationBatch;

// Abstracts the underlying data source for the High Frequency Telemetry.
// This is implemented by tlBMC and gSys.
class DataSource {
 public:
  virtual ~DataSource() = default;

  // Configures the data source for the given `identifier` with the given
  // `sampling_interval_ms`.
  // `max_buffer_size` is the maximum number of readings that the data source
  // should buffer.
  // Returns an error if the data source cannot be configured.
  virtual absl::Status ConfigureSamplingInterval(const Identifier& identifier,
                                                 int sampling_interval_ms) = 0;

  // Configures the data source for the given `identifier` with the given
  // `max_batch_size`.
  // Returns an error if the data source cannot be configured.
  virtual absl::Status ConfigureBatchSize(const Identifier& identifier,
                                          int max_batch_size) = 0;

  // Resets the data source to the default configuration for the given
  // `identifier`.
  virtual absl::Status ResetSamplingIntervalToDefault(
      const Identifier& identifier) = 0;

  // Resets the batch size to the default configuration for the given
  // `identifier`.
  virtual absl::Status ResetBatchSizeToDefault(
      const Identifier& identifier) = 0;

  // Collects the data from the data source for the given `identifier` since the
  // `start_time`.
  // Returns an error if the data source is cannot serve the data.
  virtual absl::StatusOr<Payload> Collect(const Identifier& identifier,
                                          absl::Time start_time) const = 0;

  virtual absl::StatusOr<std::vector<Identifier>> GetIdentifiersForResourceType(
      const SubscriptionPolicy::ResourceType& resource_type) const = 0;

  // Returns the name of the data source. This is used for logging and
  // debugging.
  virtual std::string GetName() const = 0;

  // Creates a sensor mutation batch for the given resource type and
  // identifiers. Only one resource type per batch.
  virtual std::unique_ptr<SensorMutationBatch> CreateSensorMutationBatch(
      SubscriptionPolicy::ResourceType resource_type,
      std::vector<Identifier> identifiers);

  // Creates a SensorMutationBatch for the given `resource_type`.
  // GetIdentifiersForResourceType will be called.
  virtual absl::StatusOr<std::unique_ptr<SensorMutationBatch>>
  CreateSensorMutationBatch(SubscriptionPolicy::ResourceType resource_type) {
    ECCLESIA_ASSIGN_OR_RETURN(auto ids,
                              GetIdentifiersForResourceType(resource_type));
    return CreateSensorMutationBatch(resource_type, std::move(ids));
  }
};

// Applies multiple sensor mutations to the data source.
// Default behavior shifts changes to Apply(), through the standard DataSource
// interface.
class SensorMutationBatch {
  // An optional sampling interval and an optional max batch size
  using Mutation = std::pair<std::optional<int>, std::optional<int>>;

 public:
  virtual ~SensorMutationBatch() = default;

  virtual absl::Status SetSamplingInterval(const Identifier& identifier,
                                           int sampling_interval_ms) & {
    ECCLESIA_ASSIGN_OR_RETURN(auto& mutation, GetMutation(identifier));
    mutation.first = sampling_interval_ms;
    return absl::OkStatus();
  }

  virtual absl::Status SetBatchSize(const Identifier& identifier,
                                    int max_batch_size) & {
    ECCLESIA_ASSIGN_OR_RETURN(auto& mutation, GetMutation(identifier));
    mutation.second = max_batch_size;
    return absl::OkStatus();
  }

  virtual absl::Status ResetSensor(const Identifier& identifier) & {
    ECCLESIA_RETURN_IF_ERROR(CheckIdentifier(identifier));
    resets_.push_back(identifier);
    return absl::OkStatus();
  }

  // Applies the mutations to the data source
  virtual absl::Status Apply() && {
    for (const auto& [identifier, mutation] : mutations_) {
      if (mutation.second.has_value()) {
        ECCLESIA_RETURN_IF_ERROR(data_source_->ConfigureBatchSize(
            identifier, mutation.second.value()));
      }
      if (mutation.first.has_value()) {
        ECCLESIA_RETURN_IF_ERROR(data_source_->ConfigureSamplingInterval(
            identifier, mutation.first.value()));
      }
    }
    for (const auto& identifier : resets_) {
      ECCLESIA_RETURN_IF_ERROR(
          data_source_->ResetSamplingIntervalToDefault(identifier));
      ECCLESIA_RETURN_IF_ERROR(
          data_source_->ResetBatchSizeToDefault(identifier));
    }
    return absl::OkStatus();
  }

 protected:
  friend class DataSource;

  SensorMutationBatch(DataSource* data_source,
                      const std::vector<Identifier>& identifiers)
      : data_source_(data_source),
        identifiers_(identifiers.begin(), identifiers.end()) {}

  // Early detection of bad identifiers prevent Apply() failures.
  absl::Status CheckIdentifier(const Identifier& identifier) {
    if (!identifiers_.contains(identifier)) {
      return absl::NotFoundError(
          absl::Substitute("Identifier $0 is not found", identifier));
    }
    return absl::OkStatus();
  }

 private:
  struct IdentifierHash {
    size_t operator()(const Identifier& v) const {
      return v.identifier_case() == Identifier::kSensorIdentifier
                 ? absl::HashOf(1, v.sensor_identifier().name())
                 : absl::HashOf(2, v.fru_identifier().barepath());
    }
  };

  struct IdentifierEqual {
    bool operator()(const Identifier& lhs, const Identifier& rhs) const {
      if (lhs.identifier_case() != rhs.identifier_case()) {
        return false;
      }
      return lhs.identifier_case() == Identifier::kSensorIdentifier
                 ? lhs.sensor_identifier().name() ==
                       rhs.sensor_identifier().name()
                 : lhs.fru_identifier().barepath() ==
                       rhs.fru_identifier().barepath();
    }
  };

  absl::StatusOr<Mutation&> GetMutation(const Identifier& identifier) {
    ECCLESIA_RETURN_IF_ERROR(CheckIdentifier(identifier));
    return mutations_[identifier];
  }

  DataSource* data_source_;
  // All sampling interval and batch size changes, per sensor identifier
  absl::flat_hash_map<Identifier, Mutation, IdentifierHash, IdentifierEqual>
      mutations_;
  std::vector<Identifier> resets_;
  absl::flat_hash_set<Identifier, IdentifierHash, IdentifierEqual> identifiers_;
};

inline std::unique_ptr<SensorMutationBatch>
DataSource::CreateSensorMutationBatch(
    SubscriptionPolicy::ResourceType /*resource_type*/,
    std::vector<Identifier> identifiers) {
  return std::unique_ptr<SensorMutationBatch>(
      new SensorMutationBatch(this, std::move(identifiers)));
}

}  // namespace milotic_hft

#endif  // THIRD_PARTY_MILOTIC_EXTERNAL_CC_HFT_ADAPTER_DATA_SOURCE_H_