package com.atlassian.messagequeue.internal.core;
   
   import cloud.atlassian.external.container.DefaultProductContainerProxy;
   import com.atlassian.messagequeue.MessageRunnerKey;
   import com.atlassian.messagequeue.internal.core.messagevalidators.MessageRunnerKeyValidator;
   import com.atlassian.messagequeue.internal.core.messagevalidators.TenantIdKeyValidator;
   import com.atlassian.messagequeue.registry.MessageContext;
   import com.atlassian.messagequeue.registry.MessageRunner;
   import com.atlassian.messagequeue.registry.MessageValidator;
  import com.atlassian.tenant.impl.TenantIdSetter;
  import com.atlassian.workcontext.api.WorkContextDoorway;
  import com.timgroup.statsd.StatsDClient;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.slf4j.MDC;
  
  import java.util.Arrays;
  import java.util.Collection;
  import java.util.List;
  import java.util.Map;
  import java.util.Optional;
  
  import static java.util.Objects.requireNonNull;
  import static java.util.concurrent.TimeUnit.MILLISECONDS;
  import static java.util.concurrent.TimeUnit.NANOSECONDS;
  
  /**
   * Consumes a {@code NestedMessage}.
   */
  public class NestedMessageConsumer {
      public static final String MDC_MESSAGE_RUNNER_KEY = "amq.messageRunnerKey";
  
      static final String STATSD_MESSAGE_PROCESSING_COUNT = "amq.message.processing.count";
      static final String STATSD_MESSAGE_PROCESSING_DURATION = "amq.message.processing.duration";
      static final String STATSD_MESSAGE_PROCESSING_TENANTID_SET = "amq.message.processing.tenantid.set";
      static final String STATSD_MSG_RUNNER_TAG_PREFIX = "msg_runner_key:";
  
      private static final Logger log = LoggerFactory.getLogger(NestedMessageConsumer.class);
  
      /**
       * The point in time that message processing completed (without error) expressed as the number of milliseconds since epoch.
       */
      private static final String MDC_MESSAGE_PROCESSING_COMPLETED_TIMESTAMP = "amq.messageProcessingCompletedTimestamp";
      private static final String MDC_MESSAGE_PROCESSING_DURATION_MILLIS = "amq.messageProcessingDurationMillis";
  
  
      private final MessageRunnerRegistryHelper messageRunnerRegistryHelper;
      private final TenantIdSetter tenantIdSetter;
      private final MessageValidatorRegistryHelper messageValidatorRegistryHelper;
  
      /**
       * Constructs an instance of {@code NestedMessageConsumer}.
       *
       * @param messageRunnerRegistryHelper    the message runner registry helper
       * @param tenantIdSetter                 the tenant ID setter
       * @param messageValidatorRegistryHelper a helper to get all message validators
       */
      public NestedMessageConsumer(MessageRunnerRegistryHelper messageRunnerRegistryHelper,
                                   TenantIdSetter tenantIdSetter,
                                   MessageValidatorRegistryHelper messageValidatorRegistryHelper) {
          this.messageRunnerRegistryHelper = requireNonNull(messageRunnerRegistryHelper);
          this.tenantIdSetter = requireNonNull(tenantIdSetter);
          this.messageValidatorRegistryHelper = requireNonNull(messageValidatorRegistryHelper);
      }
  
      private Optional<StatsDClient> getStatsDClient() {
          return DefaultProductContainerProxy.getInstance().getComponentSafely(StatsDClient.class);
      }
  
      /**
       * Consumes a {@link NestedMessage}.
       *
       * @param nestedMessage  the nested message
       * @param messageContext the message context
       */
      public void consume(NestedMessage nestedMessage, MessageContext messageContext) {
          final String messageId = messageContext.getMessageId().orElse("");
          final String tenantId = nestedMessage.getAttribute(NestedMessageConstants.TENANT_ID_ATTRIBUTE_NAME);
  
          final List<MessageValidator> defaultValidators = Arrays.asList(
                  new MessageRunnerKeyValidator(),
                  new TenantIdKeyValidator());
  
          final Map<String, String> messageAttributes = nestedMessage.getAttributesClone();
  
          for (MessageValidator validator : defaultValidators) {
              if (!validator.isValid(messageContext, messageAttributes)) {
                  validator.handleInvalidMessage(messageContext, messageAttributes);
                  return;
              }
          }
  
          try (WorkContextDoorway ignored = new WorkContextDoorway().open()) {
              tenantIdSetter.setTenantId(tenantId);
  
              final Collection<MessageValidator> validators = messageValidatorRegistryHelper.getAllValidators();
  
              for (MessageValidator validator : validators) {
                  if (!validator.isValid(messageContext, messageAttributes)) {
                     validator.handleInvalidMessage(messageContext, messageAttributes);
                     return;
                 }
             }
 
             final String messageRunnerKey = nestedMessage.getAttribute(NestedMessageConstants.MESSAGE_RUNNER_KEY_ATTRIBUTE_NAME);
             final Optional<MessageRunner> messageRunner = messageRunnerRegistryHelper.getMessageRunner(MessageRunnerKey.of(messageRunnerKey));
             if (!messageRunner.isPresent()) {
                 log.error("No MessageRunner found to process message (messageID: {}, messageRunnerKey: {})", messageId, messageRunnerKey);
                 return;
             }
 
             MDC.put(MDC_MESSAGE_RUNNER_KEY, messageRunnerKey);
 
 
             long startTimeNanos = 0;
             if (log.isInfoEnabled()) {
                 log.info("Message processing started (messageId: {}, messageRunnerKey: {})", messageId, messageRunnerKey);
                 startTimeNanos = System.nanoTime();
             }
 
             try {
                 /*
                  Copy the MDC context so we can restore it after message processing.
                  This is a defensive measure against a MessageRunner clearing the MDC context
                  */
                 final Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
                 try {
                     messageRunner.get().processMessage(messageContext);
                 } finally {
                     MDC.setContextMap(copyOfContextMap);
                 }
 
                 final long messageProcessingCompleteTimestamp = System.currentTimeMillis();
                 final long messageProcessingDurationMillis = MILLISECONDS.convert(System.nanoTime() - startTimeNanos, NANOSECONDS);
 
                 if (log.isInfoEnabled()) {
                     MDC.put(MDC_MESSAGE_PROCESSING_COMPLETED_TIMESTAMP, String.valueOf(messageProcessingCompleteTimestamp));
                     MDC.put(MDC_MESSAGE_PROCESSING_DURATION_MILLIS, String.valueOf(messageProcessingDurationMillis));
                     try {
                         log.info("Message processing completed in {}ms (messageProcessingCompletedTimestamp: {}, messageId: {}, messageRunnerKey: {})",
                                 messageProcessingDurationMillis, messageProcessingCompleteTimestamp, messageId, messageRunnerKey);
                     } finally {
                         MDC.remove(MDC_MESSAGE_PROCESSING_COMPLETED_TIMESTAMP);
                         MDC.remove(MDC_MESSAGE_PROCESSING_DURATION_MILLIS);
                     }
                 }
 
                 if (messageContext.shouldAutoAcknowledgeMessage()) {
                     messageContext.acknowledge();
                 }
 
                 postStatsDMetrics(tenantId, messageProcessingDurationMillis, STATSD_MSG_RUNNER_TAG_PREFIX + messageRunnerKey);
             } catch (Exception e) {
                 log.error("Invoking MessageRunner for message failed with an exception (messageID: {}, messageRunnerKey: {})", messageId, messageRunnerKey, e);
             }
         } finally {
             MDC.remove(MDC_MESSAGE_RUNNER_KEY);
         }
     }
 
     private void postStatsDMetrics(
             final String tenantId,
             final long messageProcessingDurationMillis,
             final String... extraStatsDTags) {
         getStatsDClient().ifPresent(statsDClient -> {
             statsDClient.increment(STATSD_MESSAGE_PROCESSING_COUNT, extraStatsDTags);
             statsDClient.recordExecutionTime(STATSD_MESSAGE_PROCESSING_DURATION, messageProcessingDurationMillis, extraStatsDTags);
             statsDClient.recordSetValue(STATSD_MESSAGE_PROCESSING_TENANTID_SET, tenantId, extraStatsDTags);
         });
     }
 }