/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * 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
 *
 *     http://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.
 */

#pragma once

#include <fb303/BaseService.h>
#include <folly/container/F14Map.h>
#include <folly/experimental/FunctionScheduler.h>
#include <thrift/lib/cpp/TProcessor.h>

namespace facebook::fb303 {

struct TStatsRequestContext {
  using clock = std::chrono::steady_clock;
  using time_point = clock::time_point;

  bool measureTime_ = false;
  bool readBeginCalled_ = false;
  bool readEndCalled_ = false;
  bool writeBeginCalled_ = false;
  bool writeEndCalled_ = false;
  bool exception = false;
  bool userException = false;
  uint32_t rBytes_ = 0;
  uint32_t wBytes_ = 0;
  // if measureTime, those will store timestamps of respective events
  time_point readBeginTime_{};
  time_point readEndTime_{};
  time_point writeBeginTime_{};
  time_point writeEndTime_{};

  void readBegin() {
    readBeginCalled_ = true;
    if (measureTime_) {
      readBeginTime_ = clock::now();
    }
  }

  void readEnd(uint32_t bytes) {
    readEndCalled_ = true;
    rBytes_ = bytes;
    if (measureTime_) {
      readEndTime_ = clock::now();
    }
  }

  void writeBegin() {
    writeBeginCalled_ = true;
    if (measureTime_) {
      writeBeginTime_ = clock::now();
    }
  }

  void writeEnd(uint32_t bytes) {
    writeEndCalled_ = true;
    wBytes_ = bytes;
    if (measureTime_) {
      writeEndTime_ = clock::now();
    }
  }

  void exceptionThrown() {
    exception = true;
  }

  void userExceptionThrown() {
    userException = true;
  }
};

/**
 * Class which holds sums, rates, counts, and times for a thrift function for
 * a single thread.  It's the smallest object of aggregation in this system.
 */
class TStatsPerThread {
 protected:
  virtual ~TStatsPerThread();

 public:
  /**
   * Get a Sampler object to return as the context for this TStatsPerThread.
   * Allocates a new Sampler if there are none left in the pool.
   */
  TStatsRequestContext* getContext();

  class StatsPerThreadHist {
   public:
    ~StatsPerThreadHist() = default;

    void setParameters(
        folly::small_vector<int> percentiles,
        CounterType bucketSize,
        CounterType min,
        CounterType max) {
      set(std::move(percentiles), bucketSize, min, max);
    }

    bool isEnabled() {
      return (exportedHist_ && hist_);
    }

    folly::small_vector<int>& getPercentiles() {
      return percentiles_;
    }

    ExportedHistogram* getExportedHistogram() {
      return exportedHist_.get();
    }

    folly::Histogram<CounterType>* getHistogram() {
      return hist_.get();
    }

   private:
    void set(
        folly::small_vector<int> percentiles,
        CounterType bucketSize,
        CounterType min,
        CounterType max);

    folly::small_vector<int> percentiles_;
    CounterType bucketSize_ = 0;
    CounterType min_ = 0;
    CounterType max_ = 0;

    std::unique_ptr<folly::Histogram<CounterType>> hist_;
    std::unique_ptr<ExportedHistogram> exportedHist_; // For consolidation
  };

  struct TimeSeries {
    uint32_t count = 0;
    uint64_t sum = 0;
    StatsPerThreadHist hist;

    void addValue(int64_t value) {
      count++;
      sum += value;
      auto h = hist.getHistogram();
      if (h) {
        h->addValue(value);
      }
    }

    void clear() {
      count = 0;
      sum = 0;
    }
  };

  /**
   * Enable histogram stats for read.
   */
  void enableThriftFuncHist(ThriftFuncHistParams* params);

  // add data from this request to stats, calling logContextDataProcessed,
  // which can be customized based on whether we observe thrift server or client
  void logContextData(const TStatsRequestContext& context);
  virtual void logContextDataProcessed(const TStatsRequestContext& context) = 0;

  void clear(); // clear all of object's sums & counts
  void setSampleRate(double rate); // set sampling fraction for timing

  std::mutex mutex_; // Mutex guarding threads stats collection
  uint32_t calls_; // total calls counted since last aggregation
  uint32_t processed_; // calls that finished processing
  uint32_t exceptions_; // total thrift undeclared exceptions counted
  uint32_t userExceptions_; // total thrift declared or undeclared exceptions
  TimeSeries readData_;
  TimeSeries writeData_;

  // timing data
  uint32_t samples_; // number of samples where timing was done
  TimeSeries readTime_;
  TimeSeries writeTime_;
  TimeSeries processTime_;
  TimeSeries totalCpuTime_;
  TimeSeries totalWorkedTime_;

  // fraction (<=1.0) of calls to be sampled for timing
  double sampleRate_ = 1.0;
  double sampleTimer_ = 0.0; // accumulates sample fractions

  // fraction of calls to be measured and logged for RequestStats
  double requestStatsMeasureRate_ = 0.0;
  double requestStatsLogRate_ = 0.0;
};

/**
 * Class which manages all function stats collection for a server.
 * To enable stats for a server, a TFunctionStatHandler should be
 * allocated,managed by a std::shared_ptr<TProcessorEventHandler>,
 * and installed in a server's TProcessor with setEventHandler().
 *
 * Only one TFunctionStatHandler should be created for a server, but
 * multiple servers within a given process space can be handled by
 * creating a TFunctionStatHandler for each.
 */
class TFunctionStatHandler
    : public apache::thrift::TProcessorEventHandler,
      public std::enable_shared_from_this<TFunctionStatHandler> {
  // we use folly::FunctionScheduler for periodic stats consolidation;
  folly::FunctionScheduler scheduler_;

  DynamicStrings dynamicStrings_; // holds the counters information
  ExportedHistogram dummyHist_;

  folly::F14NodeMap<std::string, ThriftFuncHistParams> histParamsMap_;

  /**
   * Mapping from thrift functions to their respective
   * TStatsPerThread objects for a single thread
   */
  using TStatsAggregator =
      folly::F14FastMap<std::string, std::shared_ptr<TStatsPerThread>>;

  class Tag;
  folly::ThreadLocalPtr<TStatsAggregator, Tag> tlFunctionMap_;

  std::string getHistParamsMapKey(
      std::string funcName,
      ThriftFuncAction action);

 protected:
  std::recursive_mutex statMutex_; // mutex guarding thread-local function maps

  std::string counterNamePrefix_; // "thrift.service-name."
  std::string serviceName_;
  DynamicCounters* counters_;
  int32_t nThreads_; // active threads counted last period
  int32_t secondsPerPeriod_;
  double desiredSamplesPerPeriod_; // overall samples/period wanted
  fb303::ExportedStatMap statMapSum_; // sums/rates
  fb303::ExportedStatMap statMapAvg_; // averages
  ExportedHistogramMapImpl histogramMap_; // histogram

  static const std::string kDefaultCounterNamePrefix;

  /*
   * Work to be done after the construction of TFunctionStatHandler, from the
   * destructors of child classes. This starts the
   * FunctionScheduler thread to trigger periodic consolidation.
   */
  void postConstruct();

  /*
   * Work to be done before the destruction of TFunctionStatHandler, from the
   * destructors of child classes. This triggers last consolidation and
   * shuts off the FunctionScheduler.
   * This must be done before the destructor of TFunctionStatHandler to avoid
   * a race condition on consolidate call from the scheduler.
   */
  void preDestroy();

  /*
   * This function looks for a TStatsAggregator within the current thread's
   * thread-specific memory associated with key_, creating it if necessary.
   */
  TStatsPerThread* getStats(const char* fnName);

  /**
   * Merge stats from a given thread
   */
  int32_t consolidateThread(time_t now, TStatsAggregator& functionMap);
  /**
   * Merge stats from the given TStatsPerThread object with existing counters
   */
  virtual int32_t
  consolidateStats(time_t now, const std::string& fnName, TStatsPerThread& spt);

 public:
  static const int32_t kSamplesPerSecond = 100; // default samples/second
  static const int32_t kSecondsPerPeriod = 5; // default period in seconds

  /**
   * Constructor for TFunctionStatHandler.  Initialize thread-specific
   * storage key, period timer, and hook up to dynamic counters.
   *
   * @param counters pointer to the server's dynamic counter object
   * @param sampPerSecond target # of timing samples per second
   * @param secondsPerPeriod sampling rate
   * @param useSubMinuteIntervalCounters whether or not to add high resolution
   * counters
   */
  TFunctionStatHandler(
      DynamicCounters* counters,
      const std::string& serviceName,
      int32_t sampPerSecond = kSamplesPerSecond,
      int32_t secondsPerPeriod = kSecondsPerPeriod,
      const std::string& counterNamePrefix = kDefaultCounterNamePrefix,
      bool useSubMinuteIntervalCounters = false);

  /**
   * Construct an instance of TStatsPerThread.
   */
  virtual std::shared_ptr<TStatsPerThread> createStatsPerThread(
      const char* fnName) = 0;

  /**
   * Calls setDefaultStat on all ExportedStatMapImpl members of this handler.
   * TODO(ryandm): Caveat, this function should be called before the first
   * aggregation timeout expires (~5 seconds from object creation). Task 950790.
   */
  void setDefaultStat(const ExportedStat& defaultStat) {
    statMapSum_.setDefaultStat(defaultStat);
    statMapAvg_.setDefaultStat(defaultStat);
  }

  /**
   * Aggregate the stats from all threads. Meant to be called periodically.
   */
  virtual void consolidate();

  /**
   * Get the stats-collection context for this call within the current thread.
   *
   * This obtains a thread-local TStatsAggregator and then a TStatsPerThread
   * object from that TStatsAggregator (which will create it if necessary) for
   * this function context.
   *
   * @returns current function & thread's new or existing TStatsPerThread,
   *          cast into a void*.
   */
  void* getContext(
      const char* fnName,
      apache::thrift::server::TConnectionContext* /*serverContext*/ =
          nullptr) override {
    auto stats = getStats(fnName);
    return (void*)(stats->getContext());
  }

  /**
   * Set histogram parameters for read/write/process in TStatsPerThread object.
   * This action enables histogram for the given function
   */
  void addThriftFuncHistParams(ThriftFuncHistParams params) {
    std::string key = getHistParamsMapKey(params.funcName, params.action);
    auto it = histParamsMap_.find(key);
    if (it == histParamsMap_.end()) {
      histParamsMap_.insert(std::make_pair(key, params));
    }
  }

  ThriftFuncHistParams* getThriftFuncHistParams(
      const char* fn_name,
      ThriftFuncAction action) {
    std::string key = getHistParamsMapKey(std::string(fn_name), action);
    auto it = histParamsMap_.find(key);
    if (it == histParamsMap_.end()) {
      return nullptr;
    }
    return &(it->second);
  }

  void setThriftHistParams(TStatsPerThread* stats, const char* fn_name);

  /**
   * Free resources associated with a context.
   */
  void freeContext(void* ctx, const char* fn_name) override;

  /**
   * Called before reading arguments.
   */
  void preRead(void* ctx, const char*) override;

  /**
   * Called between reading arguments and calling the handler.
   */
  void postRead(
      void* ctx,
      const char*,
      apache::thrift::transport::THeader*,
      uint32_t bytes) override;

  /**
   * Called between calling the handler and writing the response.
   */
  void preWrite(void* ctx, const char*) override;

  /**
   * Called after writing the response.
   */
  void postWrite(void* ctx, const char*, uint32_t bytes) override;

  /**
   * Called if the handler throws an undeclared exception.
   */
  void handlerError(void* ctx, const char*) override;

  /**
   * Called if the handler throws any exception.
   */
  void userExceptionWrapped(
      void* ctx,
      const char* fn_name,
      bool declared,
      const folly::exception_wrapper& ew_) final;
};

/**
 * Enable Thrift handler call counters
 */
std::shared_ptr<TFunctionStatHandler> addThriftFunctionStatHandler(
    const char* serviceName);

void withThriftFunctionStats(
    const char* serviceName,
    BaseService* service,
    folly::Function<void()>&& fn);

} // namespace facebook::fb303
