package com.atlassian.bitbucket.mesh.metrics;

import com.codahale.metrics.jmx.JmxReporter;
import com.codahale.metrics.jmx.ObjectNameFactory;
import com.google.common.base.Splitter;
import com.google.common.collect.ImmutableMap;
import io.micrometer.core.instrument.Meter;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.config.NamingConvention;
import io.micrometer.core.instrument.util.HierarchicalNameMapper;
import io.micrometer.jmx.JmxMeterRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Nonnull;
import javax.management.MalformedObjectNameException;
import javax.management.ObjectName;
import java.util.Map;
import java.util.Objects;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

import static io.micrometer.core.instrument.config.NamingConvention.camelCase;
import static java.util.Objects.requireNonNull;
import static org.apache.commons.lang3.StringUtils.isNotBlank;
import static org.apache.commons.lang3.StringUtils.stripEnd;

/**
 * Micrometer {@link HierarchicalNameMapper} and {@link ObjectNameFactory} implementation that is used to construct
 * the names of the metrics and the {@link ObjectName} used to export them over JMX.
 * <p>
 * The transformation of a micrometer {@link Meter.Id} into the {@link ObjectName} under which the metric is exposed
 * over JMX happens in a few steps in the {@link JmxMeterRegistry} support:
 * <ol>
 *     <li>when a new metric is added, {@link Meter.Id} is translated into a string using a pluggable
 *     {@link HierarchicalNameMapper}. The metric is then added to the Dropwizard {@code MetricRegistry}
 *     under that name</li>
 *     <li>this in turn triggers the {@link JmxReporter} listener to register a new
 *     {@code MBean} for the newly added metric. It uses a pluggable {@link ObjectNameFactory} to build an
 *     {@link ObjectName} from the metric <em>name</em> that was previously generated by the
 *     {@link HierarchicalNameMapper}</li>
 * </ol>
 * This class takes care of both the {@link Meter.Id} to string mapping, and the subsequent string to {@link ObjectName}
 * mapping, using both the fragments of the meter name and any {@link Tag}s provided on
 * the {@link Meter.Id} to group the metrics in cohesive groups.
 * <p>
 * In step 1, the {@link Meter.Id} is mapped to a string by taking the meter name, then a {@code /} separator and
 * appending the tags as {@code .key.value} postfixes. For instance, the {@code rpc.calls} timer with
 * {@code rpc-service=HostingService, rpc-method=HttpBackend} tags gets mapped to
 * {@code rpc.calls/rpc-service.HostingService.rpc-method.HttpBackend}
 * <p>
 * Any common tags (defined as {@code metrics.tags.*} properties) are stripped from the names. Remaining tags are
 * sorted by key using the {@link MetricsConstants#TAG_COMPARATOR} comparator, and their value appended to the
 * provided {@link Meter.Id} as described above.
 * <p>
 * In step 2, the generated string is mapped to an {@link ObjectName} by taking the metric name and interpreting each
 * dot ({@code .}) in the name as {@link ObjectName} properties (type=fragment1, section=fragment2, name=fragment3).
 * If there are more than 3 fragments in the metric name, the substring after the second dot is camelCased and used
 * as the third fragment (and thus the {@code name} property). Any tag postfixes are converted into additional
 * {@code tag-key=tag-value} {@link ObjectName} properties.
 * <p>
 * For the {@code rpc.calls} example, this results in
 * {@code metrics:type=rpc,rpc-service=HostingService,rpc-method=HttpBackend,name=calls}, which gets displayed in
 * jconsole as metric/rpc/HostingService/HttpBackend/call. For most metrics, this results in a sensible layout, but
 * sometimes we want to use one of the tags as the 'leaf' level (always the {@code name}) property. For these metrics,
 * an entry can be added to the {@link #METRIC_NAME_TAGS} map.
 * <p>
 * In {@link #METRIC_NAME_TAGS} the {@code rpc-method} tag is configured to provide the {@code name} field for
 * {@code rpc.calls}. Because of this configuration, the metric is mapped to
 * {@code metrics:type=rpc,section=calls,rpc-service=HostingService,name=HttpBackend}
 */
public class JmxRegistryNameHandler implements HierarchicalNameMapper, ObjectNameFactory {

    private static final Logger log = LoggerFactory.getLogger(JmxRegistryNameHandler.class);

    private static final Pattern PATTERN_UNSAFE = Pattern.compile("[\\s*?,=:]");
    private static final Pattern PATTERN_WORD_SEPARATOR = Pattern.compile("[\\s\\-_]+");
    //Map used to optimize the Metric.Id --> ObjectName mapping for a handful of metrics, where one of the tags should
    //be used as the JMX ObjectName name property (instead of the last fragment of the metric name). This can be used
    //for some tagged metrics. As an example, the "throttling.threads.queued" configuration ensures that the metric
    //for the "throttled-resource=scm-hosting" tag gets mapped to throttling/threads/queued/scm-hosting instead of
    //throttling/threads/scm-hosting/queued
    private static final Map<String, String> METRIC_NAME_TAGS = ImmutableMap.<String, String>builder()
            .put("jvm.gc.pause", "cause")
            .put("jvm.threads.states", "state")
            .put("logback.events", "level")
            .put("rpc.calls", "rpc-method")
            .put("rpc.messages", "rpc-method")
            .put("throttling.threads.queued", "throttled-resource")
            .build();

    private final Map<String, String> commonTags;

    public JmxRegistryNameHandler(@Nonnull Map<String, String> commonTags) {
        this.commonTags = requireNonNull(commonTags, "commonTags");
    }

    /**
     * Custom {@link ObjectNameFactory} implementation that translates the
     * {@link #toHierarchicalName(Meter.Id, NamingConvention) generated metrics name} to an {@link ObjectName}, whereby
     * related metrics are grouped together as much as possible.
     */
    @Override
    public ObjectName createName(String type, String domain, String name) {
        try {
            name = stripEnd(name, "/."); //Remove any trailing dots or slashes, just in case
            int index = name.indexOf('/');
            boolean hasTags = index != -1;
            String metricName = hasTags ? name.substring(0, index) : name;
            StringBuilder builder = new StringBuilder(domain).append(":");
            String nameProperty = addMetricNameToObjectName(builder, metricName);
            if (hasTags) {
                addMetricTagsToObjectName(builder, metricName, name.substring(index + 1));
            }
            if (nameProperty != null) {
                builder.append(",name=").append(nameProperty);
            }
            return new ObjectName(builder.toString());
        } catch (MalformedObjectNameException e) {
            try {
                return new ObjectName(domain, "name", ObjectName.quote(name));
            } catch (MalformedObjectNameException e1) {
                log.warn("Unable to register {} {}", type, name, e1);
                throw new RuntimeException(e1);
            }
        }
    }

    @Nonnull
    @Override
    public String toHierarchicalName(@Nonnull Meter.Id id, @Nonnull NamingConvention convention) {
        String tagSuffix = stripEnd(id.getTags().stream()
                //Filter out the tags that are present on all metrics. They're just noise in the JMX listing
                .filter(tag -> isNotBlank(tag.getValue()) &&
                               !Objects.equals(tag.getValue(), commonTags.get(tag.getKey())))
                //Sort the tags using the order defined in TAG_COMPARATOR
                .sorted(MetricsConstants.TAG_COMPARATOR)
                //Append the tags as .<key>.<value> suffixes
                .map(tag -> camelCase(tag.getKey()) + "." + tag.getValue().replace(".", "_") + ".")
                .map(nameSegment -> nameSegment.replace(" ", "_"))
                .collect(Collectors.joining("")), ".");
        //Separate the metric name from the tag suffix to allow it to be parsed again in the createName method.
        //Yes, it's a bit backwards to first construct a hierarchical name from the Meter.Id only to parse it again
        //before registering the JMX beans. Unfortunately, there's no simple way around it..
        return id.getConventionName(NamingConvention.dot) + (tagSuffix.isEmpty() ? "" : "/" + tagSuffix);
    }

    /**
     * Translates metric names such as {@code 'jvm.buffer.count`} to an {@link ObjectName} string where the first
     * fragment is used as the {@code type} property, the second fragment (if any) is used as the {@code section}
     * fragment and the remaining fragments are converted to camel case and used as the {@code name} property.
     * <p>
     *  Examples of translations
     *  <ul>
     *      <li>{@code jvm.buffer.count} is translated to {@code type=jvm,section=buffer,name=count}</li>
     *      <li>{@code jvm.gc.max.data.size} is translated to {@code type=jvm,section=gc,name=maxDataSize}</li>
     *      <li>{@code packCache.hits} is translated to {@code type=packCache,name=hits}</li>
     *      <li>{@code operation.git.hook.ReceivePack} is translated to
     *          {@code type=operation,operation=git,section=hook,name=ReceivePack}</li>
     *  </ul>
     *  The {@code name} property is not added to the builder yet, but returned to ensure it is added to the
     *  {@link ObjectName} string last.
     */
    private static String addMetricNameToObjectName(StringBuilder builder, String metricName) {
        int index = metricName.indexOf('.');
        String type = safeValue(index == -1 ? metricName : metricName.substring(0, index));
        builder.append("type=").append(type);
        if (index == -1 || index + 1 >= metricName.length()) {
            //Just one fragment in the metricName, we're done
            return null;
        }

        if ("operation".equals(type)) {
            //The "operation" prefix is used to group atlassian-profiling metrics. Add special handling to still
            //support 3 levels under "operation"
            int operationEnd = metricName.indexOf('.', index + 1);
            if (operationEnd != -1 && metricName.length() > operationEnd + 1) {
                //There are at least 3 parts in the metricName (type, section, name)
                builder.append(",operation=").append(safeValue(metricName.substring(index + 1, operationEnd)));
                index = operationEnd;
            }
        }

        int sectionEnd = metricName.indexOf('.', index + 1);
        if (sectionEnd != -1 && metricName.length() > sectionEnd + 1) {
            //There are at least 3 parts in the metricName (type, section, name)
            builder.append(",section=").append(safeValue(metricName.substring(index + 1, sectionEnd)));
            index = sectionEnd;
        }

        String remainder = safeValue(camelCase(metricName.substring(index + 1)));
        if (METRIC_NAME_TAGS.containsKey(metricName)) {
            //The name property will be mapped from one of the tags, add the remainder as the subsection
            builder.append(",").append(index == sectionEnd ? "subsection" : "section").append("=").append(remainder);
            return null;
        }
        return remainder;
    }

    /**
     * Translates the tag part of {@link #toHierarchicalName(Meter.Id, NamingConvention) generated hierarchical names}
     * names such as {@code 'jvm.buffer.count/id.direct`} to the substring used to build the {@link ObjectName}. The
     * tags are transformed into {@code key=value} fragments.
     */
    private static void addMetricTagsToObjectName(StringBuilder builder, String metricName, String tagsSuffix) {
        String tagUsedAsName = camelCase(METRIC_NAME_TAGS.get(metricName));
        String key = null;
        for (String part : Splitter.on('.').omitEmptyStrings().trimResults().split(tagsSuffix)) {
            if (key == null) {
                key = Objects.equals(part, tagUsedAsName) ? "name" : part;
            } else {
                builder.append(",").append(key).append("=").append(safeValue(part));
                key = null;
            }
        }
    }

    private static String camelCase(String value) {
        if (value == null) {
            return null;
        }
        //NamingConventions.camelCase only treats . as a word separator, so we convert spaces to . to treat them as
        //word separators as well.
        return camelCase.tagKey(PATTERN_WORD_SEPARATOR.matcher(value).replaceAll("."));
    }

    private static String safeValue(String value) {
        Matcher matcher = PATTERN_UNSAFE.matcher(value);
        if (matcher.find()) {
            //Unsafe character found
            return ObjectName.quote(value);
        }
        return value;
    }
}
