// 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.

#ifndef COBALT_SRC_LOGGER_PRIVACY_ENCODER_H_
#define COBALT_SRC_LOGGER_PRIVACY_ENCODER_H_

#include "src/algorithms/privacy/count_min.h"
#include "src/algorithms/random/random.h"
#include "src/lib/statusor/statusor.h"
#include "src/pb/observation.pb.h"
#include "src/registry/metric_definition.pb.h"
#include "src/registry/report_definition.pb.h"

namespace cobalt::logger {

// The PrivacyEncoder is used by the ObservationGenerator to add privacy to
// Observations (of any non-private observation type) generated by local
// aggregation procedures. The PrivateIndexObservations produced by the
// PrivacyEncoder are then written to the ObservationStore by the
// ObservationGenerator's ObservationWriter.
//
// |secure_gen| is used to perform the privacy encoding and therefore needs to be cryptographically
// secure.
// |gen| is used to perform non-deterministic encodings prior to the privacy encoding and therefore
// does not need to be cryptographically secure.
class PrivacyEncoder {
 public:
  explicit PrivacyEncoder(std::unique_ptr<SecureBitGeneratorInterface<uint32_t>> secure_gen,
                          std::unique_ptr<BitGeneratorInterface<uint32_t>> gen);

  // Returns a PrivacyEncoder with a private random number generator.
  static std::unique_ptr<PrivacyEncoder> MakeSecurePrivacyEncoder();

  virtual ~PrivacyEncoder() = default;

  // If the report definition does not specify a privacy level, MaybeMakeObservations returns
  // |observation|. Else, MaybeMakeObservations transforms a user contribution for a local
  // aggregation period into a list of private observations. If the local aggregation produced an
  // Observation for the period, then this should be passed as
  // |observation|. If not (i.e. if no events were logged during the
  // aggregation period), then |observation| should be nullptr.
  virtual lib::statusor::StatusOr<std::vector<std::unique_ptr<Observation>>>
  MaybeMakePrivateObservations(std::unique_ptr<Observation> observation,
                               const MetricDefinition &metric_def,
                               const ReportDefinition &report_def);

  // Returns the maximum index of a PrivateIndexObservation for the metric and report specified by
  // |metric_def| and |report_def|.
  static lib::statusor::StatusOr<uint64_t> MaxIndexForReport(const MetricDefinition &metric_def,
                                                             const ReportDefinition &report_def);

  // The number of sketch cells per hash function in CountMin sketch for a StringCounts report
  // |report_def|. Currently hard-coded to 10. Returns an error status if |report_def| is not of
  // type StringCounts.
  static lib::statusor::StatusOr<size_t> GetNumCountMinCellsPerHash(
      const ReportDefinition &report_def);

  // The number of hash functions in a CountMin sketch for a StringCounts report |report_def|.
  // Currently hard-coded to 5. Returns an error status if |report_def| is not of type StringCounts.
  static lib::statusor::StatusOr<size_t> GetNumCountMinHashes(const ReportDefinition &report_def);

 private:
  friend class PrivacyEncoderTest;

  lib::statusor::StatusOr<std::vector<std::unique_ptr<Observation>>> MakePrivateObservations(
      const Observation *observation, const MetricDefinition &metric_def,
      const ReportDefinition &report_def);

  lib::statusor::StatusOr<std::vector<uint64_t>> PrepareIndexVector(
      const Observation &observation, const MetricDefinition &metric_def,
      const ReportDefinition &report_def);

  static lib::statusor::StatusOr<std::vector<uint64_t>> PrepareIndexVectorForUniqueDeviceCount(
      const Observation &observation, const MetricDefinition &metric_def);

  lib::statusor::StatusOr<std::vector<uint64_t>> PrepareIndexVectorForPerDeviceIntegerReport(
      const Observation &observation, const MetricDefinition &metric_def,
      const ReportDefinition &report_def);

  lib::statusor::StatusOr<std::vector<uint64_t>> PrepareIndexVectorForFleetwideMeansReport(
      const Observation &observation, const MetricDefinition &metric_def,
      const ReportDefinition &report_def);

  static lib::statusor::StatusOr<std::vector<uint64_t>>
  PrepareIndexVectorForPerDeviceHistogramsReport(const Observation &observation,
                                                 const MetricDefinition &metric_def,
                                                 const ReportDefinition &report_def);

  lib::statusor::StatusOr<std::vector<uint64_t>> PrepareIndexVectorForFleetwideHistogramsReport(
      const Observation &observation, const MetricDefinition &metric_def,
      const ReportDefinition &report_def);

  lib::statusor::StatusOr<std::vector<uint64_t>> PrepareIndexVectorForStringCountsReport(
      const Observation &observation, const MetricDefinition &metric_def,
      const ReportDefinition &report_def, size_t num_cells_per_hash, size_t num_hashes);

  static std::vector<std::unique_ptr<Observation>> ObservationsFromIndices(
      const std::vector<uint64_t> &indices);

  lib::statusor::StatusOr<std::vector<uint64_t>> AddNoise(const std::vector<uint64_t> &indices,
                                                          const MetricDefinition &metric_def,
                                                          const ReportDefinition &report_def);

  static int64_t ClipValue(int64_t value, const ReportDefinition &report_def);

  static uint64_t ClipCount(uint64_t count, const ReportDefinition &report_def);

  std::unique_ptr<SecureBitGeneratorInterface<uint32_t>> secure_gen_;
  std::unique_ptr<BitGeneratorInterface<uint32_t>> gen_;
};

}  // namespace cobalt::logger

#endif  // COBALT_SRC_LOGGER_PRIVACY_ENCODER_H_