#ifndef PROTON_DATA_METRIC_H_
#define PROTON_DATA_METRIC_H_

#include "Runtime/Runtime.h"
#include "Utility/String.h"
#include "Utility/Traits.h"
#include <atomic>
#include <map>
#include <mutex>
#include <set>
#include <shared_mutex>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>

namespace proton {

enum class MetricKind { Flexible, Kernel, PCSampling, Cycle, Count };

using MetricValueType =
    std::variant<uint64_t, int64_t, double, std::string, std::vector<uint64_t>,
                 std::vector<int64_t>, std::vector<double>>;

inline const char *getTypeNameForIndex(std::size_t idx) {
  switch (idx) {
  case 0:
    return "uint64_t";
  case 1:
    return "int64_t";
  case 2:
    return "double";
  case 3:
    return "std::string";
  case 4:
    return "std::vector<uint64_t>";
  case 5:
    return "std::vector<int64_t>";
  case 6:
    return "std::vector<double>";
  default:
    return "<unknown>";
  }
}

/// A metric is a class that can be associated with a context.
/// `Metric` is the base class for all metrics.
/// Each `Metric` has a name and a set of values.
/// Each value could be of type `uint64_t`, `int64_t`, or `double`,
/// Each value can be inclusive (inc), exclusive (exc), or a property (pty).
/// Inclusive values are aggregated by addition and can be propagated to the
/// parent.
/// Exclusive values can be aggregated at a context but cannot be
/// propagated to the parent.
/// Property values are not aggregated and cannot be propagated to the parent.
class Metric {
public:
  Metric(MetricKind kind, size_t size) : kind(kind), values(size) {}

  virtual ~Metric() = default;

  virtual const std::string &getName() const = 0;

  virtual const std::string &getValueName(int valueId) const = 0;

  virtual bool isProperty(int valueId) const = 0;

  virtual bool isExclusive(int valueId) const = 0;

  const std::vector<MetricValueType> &getValues() const { return values; }

  const MetricValueType &getValue(int valueId) const { return values[valueId]; }

  /// Update a specific value id with the new value.
  void updateValue(int valueId, MetricValueType value) {
    // Enforce type consistency: once a valueId has a type, it must not change.
    if (values[valueId].index() != value.index()) {
      throw std::runtime_error(
          std::string("Metric value type mismatch for valueId ") +
          std::to_string(valueId) + " (" + getValueName(valueId) + ")" +
          ": current=" + getTypeNameForIndex(values[valueId].index()) +
          ", new=" + getTypeNameForIndex(value.index()));
    }
    // Handle string and other values separately
    if (std::holds_alternative<std::string>(value)) {
      values[valueId] = std::get<std::string>(value);
    } else {
      std::visit(
          [&](auto &&currentValue, auto &&otherValue) {
            using CurrentType = std::decay_t<decltype(currentValue)>;
            using ValueType = std::decay_t<decltype(otherValue)>;
            if constexpr (std::is_same_v<ValueType, CurrentType>) {
              if (isProperty(valueId)) {
                currentValue = otherValue;
              } else if constexpr (std::is_arithmetic_v<CurrentType>) {
                currentValue += otherValue;
              } else if constexpr (is_std_vector_v<CurrentType> &&
                                   std::is_arithmetic_v<
                                       typename CurrentType::value_type>) {
                if (currentValue.size() != otherValue.size()) {
                  throw std::runtime_error(
                      std::string("[PROTON] Vector metric size mismatch for "
                                  "valueId ") +
                      std::to_string(valueId) + " (" + getValueName(valueId) +
                      "): current=" + std::to_string(currentValue.size()) +
                      ", new=" + std::to_string(otherValue.size()));
                }
                for (size_t i = 0; i < currentValue.size(); ++i) {
                  currentValue[i] += otherValue[i];
                }
              } else {
                throw std::runtime_error(
                    std::string("[PROTON] Metric aggregation not supported for "
                                "valueId ") +
                    std::to_string(valueId) + " (" + getValueName(valueId) +
                    "): type=" + getTypeNameForIndex(values[valueId].index()));
              }
            }
          },
          values[valueId], value);
    }
  }

  /// Update all values of the metric with the same value.
  void updateValue(MetricValueType value) {
    for (int i = 0; i < values.size(); ++i) {
      updateValue(i, value);
    }
  }

  /// Update all values with another metric.
  void updateMetric(const Metric &other) {
    for (int i = 0; i < values.size(); ++i) {
      updateValue(i, other.values[i]);
    }
  }

  MetricKind getKind() const { return kind; }

private:
  const MetricKind kind;

protected:
  std::vector<MetricValueType> values;
};

/// A flexible metric is provided by users but not the backend profiling API.
/// Each flexible metric has a single value.
class FlexibleMetric : public Metric {
public:
  FlexibleMetric(const std::string &valueName,
                 std::variant<MetricValueType> value)
      : Metric(MetricKind::Flexible, 1), valueName(valueName) {
    this->exclusive = endWith(valueName, "(exc)");
    this->property = endWith(valueName, "(pty)");
    this->valueName = trim(replace(this->valueName, "(exc)", ""));
    this->valueName = trim(replace(this->valueName, "(pty)", ""));
    std::visit([&](auto &&v) { this->values[0] = v; }, value);
  }

  FlexibleMetric(const std::string &valueName,
                 std::variant<MetricValueType> value, bool property,
                 bool exclusive)
      : Metric(MetricKind::Flexible, 1), valueName(valueName),
        property(property), exclusive(exclusive) {
    std::visit([&](auto &&v) { this->values[0] = v; }, value);
  }

  const std::string &getName() const override { return name; }

  const std::string &getValueName(int valueId) const override {
    return valueName;
  }

  bool isProperty(int valueId) const override { return property; }

  bool isExclusive(int valueId) const override { return exclusive; }

private:
  bool property{};
  bool exclusive{};
  const static inline std::string name = "FlexibleMetric";
  std::string valueName;
};

class KernelMetric : public Metric {
public:
  enum kernelMetricKind : int {
    StartTime,
    EndTime,
    Invocations,
    Duration,
    DeviceId,
    DeviceType,
    StreamId,
    Count,
  };

  KernelMetric() : Metric(MetricKind::Kernel, kernelMetricKind::Count) {}

  KernelMetric(uint64_t startTime, uint64_t endTime, uint64_t invocations,
               uint64_t deviceId, uint64_t deviceType, uint64_t streamId)
      : KernelMetric() {
    this->values[StartTime] = startTime;
    this->values[EndTime] = endTime;
    this->values[Invocations] = invocations;
    this->values[Duration] = endTime - startTime;
    this->values[DeviceId] = deviceId;
    this->values[DeviceType] = deviceType;
    this->values[StreamId] = streamId;
  }

  const std::string &getName() const override { return name; }

  const std::string &getValueName(int valueId) const override {
    return VALUE_NAMES[valueId];
  }

  bool isProperty(int valueId) const override { return PROPERTY[valueId]; }

  bool isExclusive(int valueId) const override { return EXCLUSIVE[valueId]; }

private:
  const static inline bool PROPERTY[kernelMetricKind::Count] = {
      true, true, false, false, true, true, true};
  const static inline bool EXCLUSIVE[kernelMetricKind::Count] = {
      false, false, false, false, true, true, true};
  const static inline std::string VALUE_NAMES[kernelMetricKind::Count] = {
      "start_time (ns)", "end_time (ns)", "count",     "time (ns)",
      "device_id",       "device_type",   "stream_id",
  };
  const static inline std::string name = "KernelMetric";
};

class PCSamplingMetric : public Metric {
public:
  enum PCSamplingMetricKind : int {
    NumSamples,
    NumStalledSamples,
    StalledBranchResolving,
    StalledNoInstruction,
    StalledShortScoreboard,
    StalledWait,
    StalledLongScoreboard,
    StalledTexThrottle,
    StalledBarrier,
    StalledMembar,
    StalledIMCMiss,
    StalledMIOThrottle,
    StalledMathPipeThrottle,
    StalledDrain,
    StalledLGThrottle,
    StalledNotSelected,
    StalledMisc,
    StalledDispatchStall,
    StalledSleeping,
    StalledSelected,
    Count,
  };

  PCSamplingMetric()
      : Metric(MetricKind::PCSampling, PCSamplingMetricKind::Count) {}

  PCSamplingMetric(PCSamplingMetricKind kind, uint64_t samples,
                   uint64_t stalledSamples)
      : PCSamplingMetric() {
    this->values[kind] = stalledSamples;
    this->values[PCSamplingMetricKind::NumSamples] = samples;
    this->values[PCSamplingMetricKind::NumStalledSamples] = stalledSamples;
  }

  const std::string &getName() const override { return name; }

  const std::string &getValueName(int valueId) const override {
    return VALUE_NAMES[valueId];
  }

  bool isProperty(int valueId) const override { return false; }
  bool isExclusive(int valueId) const override { return false; }

  const static inline std::string VALUE_NAMES[PCSamplingMetricKind::Count] = {
      "num_samples",
      "num_stalled_samples",
      "stalled_branch_resolving",
      "stalled_no_instruction",
      "stalled_short_scoreboard",
      "stalled_wait",
      "stalled_long_scoreboard",
      "stalled_tex_throttle",
      "stalled_barrier",
      "stalled_membar",
      "stalled_imc_miss",
      "stalled_mio_throttle",
      "stalled_math_pipe_throttle",
      "stalled_drain",
      "stalled_lg_throttle",
      "stalled_not_Selected",
      "stalled_misc",
      "stalled_dispatch_stall",
      "stalled_sleeping",
      "stalled_selected",
  };
  const static inline std::string name = "PCSamplingMetric";
};

class CycleMetric : public Metric {
public:
  enum CycleMetricKind : int {
    StartCycle,
    EndCycle,
    Duration,
    NormalizedDuration,
    KernelId,
    KernelName,
    BlockId,
    ProcessorId,
    UnitId,
    DeviceId,
    DeviceType,
    TimeShiftCost,
    InitTime,
    PreFinalTime,
    PostFinalTime,
    Count,
  };

  CycleMetric() : Metric(MetricKind::Cycle, CycleMetricKind::Count) {}

  CycleMetric(uint64_t startCycle, uint64_t endCycle, uint64_t duration,
              double normalizedDuration, uint64_t kernelId,
              const std::string &kernelName, uint64_t blockId,
              uint64_t processorId, uint64_t unitId, uint64_t deviceId,
              uint64_t deviceType, uint64_t timeShiftCost, uint64_t initTime,
              uint64_t preFinalTime, uint64_t postFinalTime)
      : CycleMetric() {
    this->values[StartCycle] = startCycle;
    this->values[EndCycle] = endCycle;
    this->values[Duration] = duration;
    this->values[NormalizedDuration] = normalizedDuration;
    this->values[KernelId] = kernelId;
    this->values[KernelName] = kernelName;
    this->values[BlockId] = blockId;
    this->values[ProcessorId] = processorId;
    this->values[UnitId] = unitId;
    this->values[DeviceId] = deviceId;
    this->values[DeviceType] = deviceType;
    this->values[TimeShiftCost] = timeShiftCost;
    this->values[InitTime] = initTime;
    this->values[PreFinalTime] = preFinalTime;
    this->values[PostFinalTime] = postFinalTime;
  }

  const std::string &getName() const override { return name; }

  const std::string &getValueName(int valueId) const override {
    return VALUE_NAMES[valueId];
  }

  bool isProperty(int valueId) const override { return PROPERTY[valueId]; }

  bool isExclusive(int valueId) const override { return EXCLUSIVE[valueId]; }

private:
  const static inline bool PROPERTY[CycleMetricKind::Count] = {
      false, false, false, false, true,  true,  true, true,
      true,  true,  true,  true,  false, false, false};
  const static inline bool EXCLUSIVE[CycleMetricKind::Count] = {
      false, false, true, true, true,  true,  true, true,
      true,  true,  true, true, false, false, false};
  const static inline std::string VALUE_NAMES[CycleMetricKind::Count] = {
      "start_cycle", "end_cycle",      "cycles",         "normalized_cycles",
      "kernel_id",   "kernel_name",    "block_id",       "processor_id",
      "unit_id",     "device_id",      "device_type",    "time_shift_cost",
      "init_time",   "pre_final_time", "post_final_time"};
  const static inline std::string name = "CycleMetric";
};

/// Each TensorMetric represents a metric stored in a device buffer.
struct TensorMetric {
  uint8_t *ptr{};     // device pointer
  size_t typeIndex{}; // MetricValueType variant index
  uint64_t size{1};   // number of uint64 words stored at ptr
};

struct MetricKernelLaunchConfig {
  void *kernel{nullptr};
  unsigned int numThreads{1};
  unsigned int sharedMemBytes{0};
};

struct MetricKernelLaunchState {
  MetricKernelLaunchConfig tensor{};
  MetricKernelLaunchConfig scalar{};
  void *stream{nullptr};
};

/// Collect tensor metrics from device to host.
std::map<std::string, MetricValueType>
collectTensorMetrics(Runtime *runtime,
                     const std::map<std::string, TensorMetric> &tensorMetrics,
                     void *stream);

/// A MetricBuffer stores tensor metrics generated by GPU kernels.
/// The synchronization behaviors are handled by the runtime of the device.
/// A kernel can be associated with multiple tensor metrics but we do not
/// store the association on the device side.
///
/// Here's the layout of the buffer and it's meta data that are maintained on
/// the host:
///
///  host ->                             -------- kernel0 --------
///                                     /                         \
/// [device0] -> metric buffer -> {metric_id, value, metric_id, value, ...}
///                   |                            /|\
///                   |                             |
///                   | deviceOffsetPtr -------------
///                   | devicePtr
class MetricBuffer {
public:
  struct MetricDescriptor {
    size_t id{};
    size_t typeIndex{};
    size_t size{};
    std::string name{};
  };

public:
  MetricBuffer(size_t capacity, Runtime *runtime, bool mappedHostBuffer = true)
      : capacity(capacity), runtime(runtime),
        mappedHostBuffer(mappedHostBuffer) {}

  ~MetricBuffer();

  void receive(const std::map<std::string, TensorMetric> &tensorMetrics,
               const std::map<std::string, MetricValueType> &scalarMetrics,
               const MetricKernelLaunchState &metricKernelLaunchState);

  void reserve() { getOrCreateBuffer(); }

  Runtime *getRuntime() const { return runtime; }

  // no sync flush
  template <typename Func> void peek(Device *device, Func callback) {
    std::lock_guard<std::mutex> lock(bufferMutex);
    auto it = deviceBuffers.find(device);
    if (it != deviceBuffers.end()) {
      auto &buffer = it->second;
      callback(buffer.hostPtr);
    }
  }

  template <typename Func> void flush(Func callback, bool flushAll = false) {
    std::vector<std::pair<void *, DeviceBuffer>> buffersToFlush;
    if (flushAll) {
      std::lock_guard<std::mutex> lock(bufferMutex);
      for (auto &[device, buffer] : deviceBuffers) {
        buffersToFlush.emplace_back(device, buffer);
      }
    } else {
      buffersToFlush.emplace_back(runtime->getDevice(), getOrCreateBuffer());
    }
    for (auto &[device, buffer] : buffersToFlush) {
      synchronize(buffer);
      callback(device, buffer.hostPtr);
    }
  }

  size_t getCapacity() const { return capacity; }

  MetricDescriptor &getMetricDescriptor(size_t id) {
    std::shared_lock<std::shared_mutex> lock(metricDescriptorMutex);
    auto it = metricDescriptors.find(id);
    if (it == metricDescriptors.end()) {
      throw std::runtime_error("[PROTON] MetricBuffer: unknown metric id: " +
                               std::to_string(id));
    }
    return it->second;
  }

private:
  struct DeviceBuffer {
    uint8_t *devicePtr{};
    uint8_t *deviceOffsetPtr{};
    uint8_t *hostPtr{};
    uint64_t *hostOffset{};
    void *priorityStream{};
  };

  DeviceBuffer &getOrCreateBuffer();

  void queue(size_t metricId, TensorMetric tensorMetric, void *stream,
             const MetricKernelLaunchConfig &launchConfig);

  void queue(size_t metricId, MetricValueType scalarMetric, void *stream,
             const MetricKernelLaunchConfig &launchConfig);

  void synchronize(DeviceBuffer &buffer);

  template <typename MetricT>
  size_t getMetricTypeIndex(const MetricT &metric) const {
    using MetricType = std::decay_t<MetricT>;
    if constexpr (std::is_same_v<MetricValueType, MetricType>) {
      return metric.index();
    } else if constexpr (std::is_same_v<TensorMetric, MetricType>) {
      return metric.typeIndex;
    } else {
      static_assert(always_false<MetricType>::value,
                    "Unsupported metric type for getMetricTypeIndex");
    }
  }

  template <typename MetricT>
  size_t getMetricSize(const MetricT &metric) const {
    using MetricType = std::decay_t<MetricT>;
    if constexpr (std::is_same_v<MetricValueType, MetricType>) {
      return 1; // FIXME: We don't queue a scalar metric with vector type yet
    } else if constexpr (std::is_same_v<TensorMetric, MetricType>) {
      return metric.size;
    } else {
      static_assert(always_false<MetricType>::value,
                    "Unsupported metric type for getMetricSize");
    }
  }

  template <typename MetricsT>
  void queueMetrics(const MetricsT &metrics, void *stream,
                    const MetricKernelLaunchConfig &launchConfig) {
    for (const auto &[name, metric] : metrics) {
      size_t typeIndex = getMetricTypeIndex(metric);
      size_t size = getMetricSize(metric);
      auto descriptor = getOrCreateMetricDescriptor(name, typeIndex, size);
      queue(descriptor.id, metric, stream, launchConfig);
    }
  }

  MetricDescriptor getOrCreateMetricDescriptor(const std::string &name,
                                               size_t typeIndex, size_t size);

protected:
  static std::atomic<size_t> metricId;
  static std::map<size_t, MetricDescriptor> metricDescriptors;
  static std::map<std::string, size_t> metricNameToId;
  static std::shared_mutex metricDescriptorMutex;

  size_t capacity; // byte
  Runtime *runtime{};
  const bool mappedHostBuffer{true};

  std::map<void *, DeviceBuffer> deviceBuffers;
  std::mutex bufferMutex;
};

class MetricInterface {
public:
  virtual ~MetricInterface() = default;

  virtual void
  addMetrics(size_t scopeId,
             const std::map<std::string, MetricValueType> &scalarMetrics,
             const std::map<std::string, TensorMetric> &tensorMetrics) = 0;

  virtual void
  setMetricKernels(const MetricKernelLaunchState &metricKernelLaunchState) = 0;
};

} // namespace proton

#endif // PROTON_DATA_METRIC_H_
