package io.ray.runtime.metric;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** MetricRegistry is a registry for metrics to be registered and updates metrics. */
public class MetricRegistry {

  public static final MetricRegistry DEFAULT_REGISTRY = new MetricRegistry();

  private static final Logger LOG = LoggerFactory.getLogger(MetricRegistry.class);

  private final Map<MetricId, Metric> registeredMetrics = new HashMap<>(64);

  private MetricConfig metricConfig;
  private ScheduledExecutorService scheduledExecutorService;
  private volatile boolean isRunning = false;

  public void startup() {
    startup(MetricConfig.DEFAULT_CONFIG);
  }

  public void startup(MetricConfig metricConfig) {
    synchronized (this) {
      if (!isRunning) {
        this.metricConfig = metricConfig;
        scheduledExecutorService =
            new ScheduledThreadPoolExecutor(
                metricConfig.threadPoolSize(),
                new ThreadFactoryBuilder().setNameFormat("metric-registry-%d").build());
        scheduledExecutorService.scheduleAtFixedRate(
            this::update,
            metricConfig.timeIntervalMs(),
            metricConfig.timeIntervalMs(),
            TimeUnit.MILLISECONDS);
        isRunning = true;
        LOG.info("Finished startup metric registry, metricConfig is {}.", metricConfig);
      }
    }
  }

  public void shutdown() {
    synchronized (this) {
      if (isRunning && scheduledExecutorService != null) {
        try {
          scheduledExecutorService.shutdownNow();
          if (!scheduledExecutorService.awaitTermination(
              metricConfig.shutdownWaitTimeMs(), TimeUnit.MILLISECONDS)) {
            LOG.warn(
                "Metric registry did not shut down in {}ms time, so try to shut down again.",
                metricConfig.shutdownWaitTimeMs());
            scheduledExecutorService.shutdownNow();
          }
        } catch (InterruptedException e) {
          LOG.warn(
              "Interrupted when shutting down metric registry, so try to shut down again.",
              e.getMessage(),
              e);
          scheduledExecutorService.shutdownNow();
        }
        if (scheduledExecutorService.isShutdown()) {
          isRunning = false;
          scheduledExecutorService = null;
          LOG.info("Metric registry has been shut down.");
        } else {
          LOG.warn("Failed to shut down metric registry service.");
        }
      }
    }
  }

  public Metric register(Metric metric) {
    synchronized (this) {
      if (!isRunning) {
        LOG.warn("Failed to register a metric, because the metric registry is not running.");
        return null;
      }
      try {
        MetricId id = genMetricIdByMetric(metric);
        Metric ori = registeredMetrics.putIfAbsent(id, metric);
        if (ori == null) {
          return metric;
        } else {
          LOG.info("Metric {} has already registered, so use the previous one.", id);
          return ori;
        }
      } catch (Exception e) {
        LOG.warn("Failed to register a metric: ", e.getMessage(), e);
        return null;
      }
    }
  }

  public void unregister(Metric metric) {
    synchronized (this) {
      if (!isRunning) {
        LOG.warn("Failed to unregister a metric, because the metric registry is not running.");
      }
      try {
        MetricId id = genMetricIdByMetric(metric);
        registeredMetrics.remove(id);
      } catch (Exception e) {
        LOG.warn("Failed to unregister a metric: ", e.getMessage(), e);
      }
    }
  }

  private void update() {
    registeredMetrics.forEach(
        (id, metric) -> {
          metric.record();
        });
  }

  private MetricType getMetricType(Metric metric) {
    if (metric instanceof Count) {
      return MetricType.COUNT;
    }
    if (metric instanceof Gauge) {
      return MetricType.GAUGE;
    }
    if (metric instanceof Sum) {
      return MetricType.SUM;
    }
    if (metric instanceof Histogram) {
      return MetricType.HISTOGRAM;
    }
    throw new RuntimeException(
        "Unknown metric type, the metric is " + metric.getClass().getSimpleName());
  }

  private MetricId genMetricIdByMetric(Metric metric) {
    return new MetricId(getMetricType(metric), metric.name, metric.tags);
  }
}
