package io.prometheus.client.hibernate;

import io.prometheus.client.Collector;
import io.prometheus.client.CollectorRegistry;
import io.prometheus.client.CounterMetricFamily;
import io.prometheus.client.GaugeMetricFamily;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;
import org.hibernate.SessionFactory;
import org.hibernate.stat.Statistics;

/**
 * Collect metrics from one or more Hibernate SessionFactory instances.
 *
 * <p>Usage example for a single session factory:
 *
 * <pre>
 * new HibernateStatisticsCollector(sessionFactory, "myapp").register();
 * </pre>
 *
 * Usage example for multiple session factories:
 *
 * <pre>
 * new HibernateStatisticsCollector()
 *     .add(sessionFactory1, "myapp1")
 *     .add(sessionFactory2, "myapp2")
 *     .register();
 * </pre>
 *
 * If you are in a JPA environment, you can obtain the SessionFactory like this:
 *
 * <pre>
 * SessionFactory sessionFactory =
 *     entityManagerFactory.unwrap(SessionFactory.class);
 * </pre>
 *
 * <p>When {@code enablePerQueryMetrics()} has been called, certain metrics like execution time are
 * collected per query. This may create a lot of monitoring data, so it should be used with caution.
 *
 * @author Christian Kaltepoth
 */
class HibernaTestatisticsCollector extends Collector {

  private static final List<String> LABEL_NAMES = Collections.singletonList("unit");

  private static final List<String> LABEL_NAMES_PER_QUERY = Arrays.asList("unit", "query");

  private final Map<String, SessionFactory> sessionFactories =
      new ConcurrentHashMap<String, SessionFactory>();

  private boolean perQueryMetricsEnabled;

  /**
   * Creates an empty collector. If you use this constructor, you have to add one or more session
   * factories to the collector by calling the {@link #add(SessionFactory, String)} method.
   */
  public HibernateStatisticsCollector() {
    // nothing
  }

  /**
   * Creates a new collector for the given session factory. Calling this constructor has the same
   * effect as creating an empty collector and adding the session factory using {@link
   * #add(SessionFactory, String)}.
   *
   * @param sessionFactory The Hibernate SessionFactory to collect metrics for
   * @param name A unique name for this SessionFactory
   */
  public HibernateStatisticsCollector(SessionFactory sessionFactory, String name) {
    add(sessionFactory, name);
  }

  /**
   * Registers a Hibernate SessionFactory with this collector.
   *
   * @param sessionFactory The Hibernate SessionFactory to collect metrics for
   * @param name A unique name for this SessionFactory
   * @return Returns the collector
   */
  public HibernateStatisticsCollector add(SessionFactory sessionFactory, String name) {
    sessionFactories.put(name, sessionFactory);
    return this;
  }

  /**
   * Enables collection of per-query metrics. Produces a lot of monitoring data, so use with
   * caution.
   *
   * <p>Per-query metrics have a label "query" with the actual HQL query as value. The query will
   * contain placeholders ("?") instead of the real parameter values (example: {@code select u from
   * User u where id=?}).
   *
   * @return Returns the collector
   */
  public HibernateStatisticsCollector enablePerQueryMetrics() {
    this.perQueryMetricsEnabled = true;
    return this;
  }

  @Override
  public List<MetricFamilySamples> collect() {
    List<MetricFamilySamples> metrics = new ArrayList<MetricFamilySamples>();
    metrics.addAll(getSessionMetrics());
    metrics.addAll(getConnectionMetrics());
    metrics.addAll(getCacheMetrics());
    metrics.addAll(getEntityMetrics());
    metrics.addAll(getQueryExecutionMetrics());
    if (perQueryMetricsEnabled) {
      metrics.addAll(getPerQueryMetrics());
    }
    return metrics;
  }

  @Override
  public <T extends Collector> T register(CollectorRegistry registry) {
    if (sessionFactories.isEmpty()) {
      throw new IllegalStateException("You must register at least one SessionFactory.");
    }
    return super.register(registry);
  }

  private List<MetricFamilySamples> getSessionMetrics() {
    return Arrays.<MetricFamilySamples>asList(
        createCounter(
            "hibernate_session_opened_total",
            "Global number of sessions opened (getSessionOpenCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getSessionOpenCount();
              }
            }),
        createCounter(
            "hibernate_session_closed_total",
            "Global number of sessions closed (getSessionCloseCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getSessionCloseCount();
              }
            }),
        createCounter(
            "hibernate_flushed_total",
            "The global number of flushes executed by sessions (getFlushCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getFlushCount();
              }
            }),
        createCounter(
            "hibernate_connect_total",
            "The global number of connections requested by the sessions (getConnectCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getConnectCount();
              }
            }),
        createCounter(
            "hibernate_optimistic_failure_total",
            "The number of StaleObjectStateExceptions that occurred (getOptimisticFailureCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getOptimisticFailureCount();
              }
            }));
  }

  private List<MetricFamilySamples> getConnectionMetrics() {
    return Arrays.<MetricFamilySamples>asList(
        createCounter(
            "hibernate_statement_prepared_total",
            "The number of prepared statements that were acquired (getPrepareStatementCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getPrepareStatementCount();
              }
            }),
        createCounter(
            "hibernate_statement_closed_total",
            "The number of prepared statements that were released (getCloseStatementCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getCloseStatementCount();
              }
            }),
        createCounter(
            "hibernate_transaction_total",
            "The number of transactions we know to have completed (getTransactionCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getTransactionCount();
              }
            }),
        createCounter(
            "hibernate_transaction_success_total",
            "The number of transactions we know to have been successful"
                + " (getSuccessfulTransactionCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getSuccessfulTransactionCount();
              }
            }));
  }

  private List<MetricFamilySamples> getCacheMetrics() {
    return Arrays.<MetricFamilySamples>asList(
        createCounter(
            "hibernate_second_level_cache_hit_total",
            "Global number of cacheable entities/collections successfully retrieved from the cache"
                + " (getSecondLevelCacheHitCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getSecondLevelCacheHitCount();
              }
            }),
        createCounter(
            "hibernate_second_level_cache_miss_total",
            "Global number of cacheable entities/collections not found in the cache and loaded from"
                + " the database (getSecondLevelCacheMissCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getSecondLevelCacheMissCount();
              }
            }),
        createCounter(
            "hibernate_second_level_cache_put_total",
            "Global number of cacheable entities/collections put in the cache"
                + " (getSecondLevelCachePutCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getSecondLevelCachePutCount();
              }
            }),
        createCounter(
            "hibernate_query_cache_hit_total",
            "The global number of cached queries successfully retrieved from cache"
                + " (getQueryCacheHitCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getQueryCacheHitCount();
              }
            }),
        createCounter(
            "hibernate_query_cache_miss_total",
            "The global number of cached queries not found in cache (getQueryCacheMissCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getQueryCacheMissCount();
              }
            }),
        createCounter(
            "hibernate_query_cache_put_total",
            "The global number of cacheable queries put in cache (getQueryCachePutCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getQueryCachePutCount();
              }
            }),
        createCounter(
            "hibernate_natural_id_cache_hit_total",
            "The global number of cached naturalId lookups successfully retrieved from cache"
                + " (getNaturalIdCacheHitCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getNaturalIdCacheHitCount();
              }
            }),
        createCounter(
            "hibernate_natural_id_cache_miss_total",
            "The global number of cached naturalId lookups not found in cache"
                + " (getNaturalIdCacheMissCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getNaturalIdCacheMissCount();
              }
            }),
        createCounter(
            "hibernate_natural_id_cache_put_total",
            "The global number of cacheable naturalId lookups put in cache"
                + " (getNaturalIdCachePutCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getNaturalIdCachePutCount();
              }
            }),
        createCounter(
            "hibernate_update_timestamps_cache_hit_total",
            "The global number of timestamps successfully retrieved from cache"
                + " (getUpdateTimestampsCacheHitCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getUpdateTimestampsCacheHitCount();
              }
            }),
        createCounter(
            "hibernate_update_timestamps_cache_miss_total",
            "The global number of tables for which no update timestamps was not found in cache"
                + " (getUpdateTimestampsCacheMissCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getUpdateTimestampsCacheMissCount();
              }
            }),
        createCounter(
            "hibernate_update_timestamps_cache_put_total",
            "The global number of timestamps put in cache (getUpdateTimestampsCachePutCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getUpdateTimestampsCachePutCount();
              }
            }));
  }

  private List<MetricFamilySamples> getEntityMetrics() {
    return Arrays.<MetricFamilySamples>asList(
        createCounter(
            "hibernate_entity_delete_total",
            "Global number of entity deletes (getEntityDeleteCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getEntityDeleteCount();
              }
            }),
        createCounter(
            "hibernate_entity_insert_total",
            "Global number of entity inserts (getEntityInsertCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getEntityInsertCount();
              }
            }),
        createCounter(
            "hibernate_entity_load_total",
            "Global number of entity loads (getEntityLoadCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getEntityLoadCount();
              }
            }),
        createCounter(
            "hibernate_entity_fetch_total",
            "Global number of entity fetches (getEntityFetchCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getEntityFetchCount();
              }
            }),
        createCounter(
            "hibernate_entity_update_total",
            "Global number of entity updates (getEntityUpdateCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getEntityUpdateCount();
              }
            }),
        createCounter(
            "hibernate_collection_load_total",
            "Global number of collections loaded (getCollectionLoadCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getCollectionLoadCount();
              }
            }),
        createCounter(
            "hibernate_collection_fetch_total",
            "Global number of collections fetched (getCollectionFetchCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getCollectionFetchCount();
              }
            }),
        createCounter(
            "hibernate_collection_update_total",
            "Global number of collections updated (getCollectionUpdateCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getCollectionUpdateCount();
              }
            }),
        createCounter(
            "hibernate_collection_remove_total",
            "Global number of collections removed (getCollectionRemoveCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getCollectionRemoveCount();
              }
            }),
        createCounter(
            "hibernate_collection_recreate_total",
            "Global number of collections recreated (getCollectionRecreateCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getCollectionRecreateCount();
              }
            }));
  }

  private List<MetricFamilySamples> getQueryExecutionMetrics() {
    return Arrays.<MetricFamilySamples>asList(
        createCounter(
            "hibernate_query_execution_total",
            "Global number of executed queries (getQueryExecutionCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getQueryExecutionCount();
              }
            }),
        createCounter(
            "hibernate_natural_id_query_execution_total",
            "The global number of naturalId queries executed against the database"
                + " (getNaturalIdQueryExecutionCount)",
            new ValueProvider() {
              @Override
              public double getValue(Statistics statistics) {
                return statistics.getNaturalIdQueryExecutionCount();
              }
            }));
  }

  private List<MetricFamilySamples> getPerQueryMetrics() {
    List<MetricFamilySamples> metrics = new ArrayList<MetricFamilySamples>();

    metrics.addAll(
        Arrays.asList(
            createCounterForQuery(
                "hibernate_per_query_cache_hit_total",
                "Global number of cache hits for query (getCacheHitCount)",
                new ValueProviderPerQuery() {
                  @Override
                  public double getValue(Statistics statistics, String query) {
                    return statistics.getQueryStatistics(query).getCacheHitCount();
                  }
                }),
            createCounterForQuery(
                "hibernate_per_query_cache_miss_total",
                "Global number of cache misses for query (getCacheMissCount)",
                new ValueProviderPerQuery() {
                  @Override
                  public double getValue(Statistics statistics, String query) {
                    return statistics.getQueryStatistics(query).getCacheMissCount();
                  }
                }),
            createCounterForQuery(
                "hibernate_per_query_cache_put_total",
                "Global number of cache puts for query (getCachePutCount)",
                new ValueProviderPerQuery() {
                  @Override
                  public double getValue(Statistics statistics, String query) {
                    return statistics.getQueryStatistics(query).getCachePutCount();
                  }
                }),
            createCounterForQuery(
                "hibernate_per_query_execution_total",
                "Global number of executions for query (getExecutionCount)",
                new ValueProviderPerQuery() {
                  @Override
                  public double getValue(Statistics statistics, String query) {
                    return statistics.getQueryStatistics(query).getExecutionCount();
                  }
                }),
            createCounterForQuery(
                "hibernate_per_query_execution_rows_total",
                "Global number of rows for all executions of query (getExecutionRowCount)",
                new ValueProviderPerQuery() {
                  @Override
                  public double getValue(Statistics statistics, String query) {
                    return statistics.getQueryStatistics(query).getExecutionRowCount();
                  }
                }),
            createGaugeForQuery(
                "hibernate_per_query_execution_min_seconds",
                "Minimum execution time of query in seconds (based on getExecutionMinTime)",
                new ValueProviderPerQuery() {
                  @Override
                  public double getValue(Statistics statistics, String query) {
                    return toSeconds(statistics.getQueryStatistics(query).getExecutionMinTime());
                  }
                }),
            createGaugeForQuery(
                "hibernate_per_query_execution_max_seconds",
                "Maximum execution time of query in seconds (based on getExecutionMaxTime)",
                new ValueProviderPerQuery() {
                  @Override
                  public double getValue(Statistics statistics, String query) {
                    return toSeconds(statistics.getQueryStatistics(query).getExecutionMaxTime());
                  }
                }),
            createCounterForQuery(
                "hibernate_per_query_execution_seconds_total",
                "Accumulated execution time of query in seconds (based on getExecutionTotalTime)",
                new ValueProviderPerQuery() {
                  @Override
                  public double getValue(Statistics statistics, String query) {
                    return toSeconds(statistics.getQueryStatistics(query).getExecutionTotalTime());
                  }
                })));

    return metrics;
  }

  private CounterMetricFamily createCounter(String metric, String help, ValueProvider provider) {

    CounterMetricFamily metricFamily = new CounterMetricFamily(metric, help, LABEL_NAMES);

    for (Entry<String, SessionFactory> entry : sessionFactories.entrySet()) {
      metricFamily.addMetric(
          Collections.singletonList(entry.getKey()),
          provider.getValue(entry.getValue().getStatistics()));
    }

    return metricFamily;
  }

  private CounterMetricFamily createCounterForQuery(
      String metric, String help, ValueProviderPerQuery provider) {

    final CounterMetricFamily counters =
        new CounterMetricFamily(metric, help, LABEL_NAMES_PER_QUERY);

    addMetricsForQuery(
        new PerQuerySamples() {
          @Override
          public void addMetric(List<String> labelValues, double value) {
            counters.addMetric(labelValues, value);
          }
        },
        provider);

    return counters;
  }

  private GaugeMetricFamily createGaugeForQuery(
      String metric, String help, ValueProviderPerQuery provider) {

    final GaugeMetricFamily gauges = new GaugeMetricFamily(metric, help, LABEL_NAMES_PER_QUERY);

    addMetricsForQuery(
        new PerQuerySamples() {
          @Override
          public void addMetric(List<String> labelValues, double value) {
            gauges.addMetric(labelValues, value);
          }
        },
        provider);

    return gauges;
  }

  private void addMetricsForQuery(PerQuerySamples samples, ValueProviderPerQuery provider) {

    for (Entry<String, SessionFactory> entry : sessionFactories.entrySet()) {
      SessionFactory sessionFactory = entry.getValue();
      Statistics stats = sessionFactory.getStatistics();
      String unitName = entry.getKey();

      for (String query : stats.getQueries()) {
        samples.addMetric(Arrays.asList(unitName, query), provider.getValue(stats, query));
      }
    }
  }

  private double toSeconds(long milliseconds) {
    return milliseconds / 1000d;
  }

  private interface PerQuerySamples {

    void addMetric(List<String> labelValues, double value);
  }

  private interface ValueProvider {

    double getValue(Statistics statistics);
  }

  private interface ValueProviderPerQuery {

    double getValue(Statistics statistics, String query);
  }
}
