package com.atlassian.messagequeue.internal.scheduler;
   
   import com.atlassian.messagequeue.Message;
   import com.atlassian.messagequeue.MessageInformationService;
   import com.atlassian.messagequeue.MessageRunnerKey;
   import com.atlassian.messagequeue.MessageRunnerNotRegisteredException;
   import com.atlassian.messagequeue.MessageRunnerService;
   import com.atlassian.messagequeue.MessageRunnerServiceException;
   import com.atlassian.messagequeue.internal.core.MessageRunnerRegistryHelper;
  import com.atlassian.messagequeue.internal.core.NestedMessage;
  import com.atlassian.messagequeue.internal.core.NestedMessageConstants;
  import com.atlassian.messagequeue.internal.core.NestedMessageSerializer;
  import com.atlassian.messagequeue.registry.MessageContext;
  import com.atlassian.messagequeue.registry.MessageRunner;
  import com.atlassian.scheduler.JobRunner;
  import com.atlassian.scheduler.JobRunnerRequest;
  import com.atlassian.scheduler.JobRunnerResponse;
  import com.atlassian.scheduler.SchedulerService;
  import com.atlassian.scheduler.SchedulerServiceException;
  import com.atlassian.scheduler.config.JobConfig;
  import com.atlassian.scheduler.config.JobId;
  import com.atlassian.scheduler.config.JobRunnerKey;
  import com.atlassian.scheduler.config.RunMode;
  import com.atlassian.scheduler.config.Schedule;
  import com.atlassian.workcontext.api.WorkContextAvailable;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import javax.annotation.Nullable;
  import javax.annotation.PostConstruct;
  import javax.annotation.PreDestroy;
  
  import java.io.Serializable;
  import java.util.Date;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.Optional;
  import java.util.concurrent.TimeUnit;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * An implementation of {@code MessageRunnerService} that is backed by the {@link SchedulerService} that provides
   * at-least-once delivery semantics.
   *
   * <p>When a {@code Message} is passed to {@link #addMessage(Message)}, a job is scheduled to run
   * as-soon-as-possible to process this {@code Message}.
   *
   * <p>When the processing of a {@code Message} fails (i.e. {@link MessageRunner#processMessage(MessageContext)} throws an exception),
   * the {@code Message} will be delivered up to a total of {@code maxDeliveryCount} times at intervals of {@code deliveryIntervalMillis}
   * milliseconds until it is processed successfully. The values of these parameters can be specified when constructing an
   * instance of this class.
   *
   * <p>The processing time of a {@code Message} can exceed {@code deliveryIntervalMillis}. When this happens, re-delivery
   * of the {@code Message} is skipped. Another re-delivery will be attempted after {@code deliveryIntervalMillis} milliseconds.
   * Therefore, only one {@code MessageRunner} will be processing a particular {@code Message} at any given time.
   *
   * <p>The processing time of a {@code Message} is not currently limited. However long processing times are discouraged.
   *
   * <p>If a {@code Message} fails to be processed successfully after {@code maxDeliveryCount} attempts,
   * a {@link com.atlassian.scheduler.status.RunDetails} instance is recorded to note this occurence. A collection of
   * {@link com.atlassian.scheduler.status.RunDetails} can be viewed as a dead letter queue.
   *
   * <p>When constructing an instance of this class with a {@link SchedulerService} implementation that supports
   * persistence of jobs, all {@code Message}s passed to {@link #addMessage(Message)} will be persisted.
   *
   * <p>To properly initialise an instance of this service, a client must call {@link #init()} after
   * constructing an instance (or else ensure the dependency injection framework will automatically invoke methods
   * annotated with @PostConstruct).
   *
   * <p>To properly dispose of an instance of this service, a client must call {@link #shutdown()} (or else ensure the
   * dependency injection framework will automatically invoke methods annotated with @PreDestroy).
   */
  public class SchedulerMessageRunnerService implements MessageRunnerService {
      static final String DELIVERY_COUNT_PARAMETER = "deliveryCount";
      static final JobRunnerKey AMQ_JOB_RUNNER_KEY = JobRunnerKey.of("amqJobRunner");
  
      private static final Logger log = LoggerFactory.getLogger(SchedulerMessageRunnerService.class);
      private static final long DEFAULT_DELIVERY_INTERVAL_MILLIS = TimeUnit.SECONDS.toMillis(Integer.getInteger("amq.scheduler.delivery.interval.seconds", 30));
      private static final int DEFAULT_MAX_DELIVERY_COUNT = Integer.getInteger("amq.scheduler.max.delivery.count", 3);
      private static final String PAYLOAD_JOB_PARAMETER = "payload";
  
      private final SchedulerService schedulerService;
      private final MessageRunnerRegistryHelper messageRunnerRegistryHelper;
      private final NestedMessageSerializer nestedMessageSerializer;
      private final MessageInformationService messageInformationService;
      private final long deliveryIntervalMillis;
      private final int maxDeliveryCount;
      private volatile boolean shutdown = false;
  
      /**
       * Constructs a new instance of SchedulerMessageRunnerService.
       * @param schedulerService the scheduler service
       * @param messageRunnerRegistryHelper the message runner registry helper
       * @param messageInformationService message information service
       * @param nestedMessageSerializer the nested message serializer
       */
      public SchedulerMessageRunnerService(SchedulerService schedulerService,
                                           MessageRunnerRegistryHelper messageRunnerRegistryHelper,
                                          MessageInformationService messageInformationService,
                                          NestedMessageSerializer nestedMessageSerializer) {
         this(schedulerService, messageRunnerRegistryHelper, messageInformationService, nestedMessageSerializer,
                 DEFAULT_DELIVERY_INTERVAL_MILLIS, DEFAULT_MAX_DELIVERY_COUNT);
     }
 
     /**
      * Constructs a new instance of SchedulerMessageRunnerService.
      * @param schedulerService the scheduler service
      * @param messageRunnerRegistryHelper the message runner registry helper
      * @param messageInformationService message information service
      * @param nestedMessageSerializer nested message serializer
      * @param deliveryIntervalMillis the time interval in milliseconds between successive deliveries of a message (must be greater than 0)
      * @param maxDeliveryCount the maximum number of times to deliver a message if message processing fails and the message is not acknowledged (must be greater than 0)
      */
     public SchedulerMessageRunnerService(SchedulerService schedulerService,
                                          MessageRunnerRegistryHelper messageRunnerRegistryHelper,
                                          MessageInformationService messageInformationService, NestedMessageSerializer nestedMessageSerializer,
                                          long deliveryIntervalMillis, int maxDeliveryCount) {
         this.schedulerService = requireNonNull(schedulerService, "schedulerService");
         this.messageRunnerRegistryHelper = requireNonNull(messageRunnerRegistryHelper, "messageRunnerRegistryHelper");
         this.nestedMessageSerializer = requireNonNull(nestedMessageSerializer, "nestedMessageSerializer");
         this.messageInformationService = requireNonNull(messageInformationService, "messageInformationService");
 
         if (deliveryIntervalMillis <= 0) {
             throw new IllegalArgumentException("deliveryIntervalMillis must be greater than 0. Received: " + deliveryIntervalMillis);
         }
         this.deliveryIntervalMillis = deliveryIntervalMillis;
 
         if (maxDeliveryCount <= 0) {
             throw new IllegalArgumentException("maxDeliveryCount must be greater than 0. Received: " + maxDeliveryCount);
         }
         this.maxDeliveryCount = maxDeliveryCount;
 
         log.info("Constructing {} (maxDeliveryCount: {}, deliveryIntervalMinutes: {})", this.getClass().getSimpleName(), maxDeliveryCount, TimeUnit.MILLISECONDS.toMinutes(deliveryIntervalMillis));
     }
 
     /**
      * Initialises this instance.
      */
     @PostConstruct
     public void init() {
         log.info("Registering AMQ JobRunner");
         schedulerService.registerJobRunner(AMQ_JOB_RUNNER_KEY, new AmqJobRunner());
     }
 
     private class AmqJobRunner implements JobRunner {
         @Nullable
         @Override
         public JobRunnerResponse runJob(JobRunnerRequest jobRunnerRequest) {
             final JobRunnerResponse validateResponse = validate(jobRunnerRequest);
             if (validateResponse != null) {
                 return validateResponse;
             }
 
             final JobRunnerResponse incrementDeliveryCountResponse = incrementDeliveryCount(jobRunnerRequest);
             if (incrementDeliveryCountResponse != null) {
                 return incrementDeliveryCountResponse;
             }
 
             final String payload = (String) jobRunnerRequest.getJobConfig().getParameters().get(PAYLOAD_JOB_PARAMETER);
             final NestedMessage nestedMessage = nestedMessageSerializer.deserialize(payload);
 
             return processNestedMessage(nestedMessage, new SchedulerMessageContext(jobRunnerRequest, nestedMessage.getPayload(), schedulerService));
         }
 
         JobRunnerResponse processNestedMessage(NestedMessage nestedMessage, MessageContext messageContext) {
             final String messageId = messageContext.getMessageId().orElse("");
             final String tenantId = nestedMessage.getAttribute(NestedMessageConstants.TENANT_ID_ATTRIBUTE_NAME);
             if (tenantId == null) {
                 final String errorMessageFormat = "Message received that does not contain required '%s' attribute. Message will be acknowledged (messageID: %s)";
                 final String errorMessage = String.format(errorMessageFormat, NestedMessageConstants.TENANT_ID_ATTRIBUTE_NAME, messageId);
                 log.error(errorMessage);
                 messageContext.acknowledge();
 
                 return JobRunnerResponse.aborted(errorMessage);
             }
 
             final String messageRunnerKey = nestedMessage.getAttribute(NestedMessageConstants.MESSAGE_RUNNER_KEY_ATTRIBUTE_NAME);
             if (messageRunnerKey == null) {
                 final String errorMessageFormat = "Message received that does not contain required '%s' attribute. Message will be acknowledged (messageID: %s)";
                 final String errorMessage = String.format(errorMessageFormat, NestedMessageConstants.MESSAGE_RUNNER_KEY_ATTRIBUTE_NAME, messageId);
                 log.error(errorMessage);
                 messageContext.acknowledge();
 
                 return JobRunnerResponse.aborted(errorMessage);
             }
 
             final Optional<MessageRunner> messageRunner = messageRunnerRegistryHelper.getMessageRunner(MessageRunnerKey.of(messageRunnerKey));
             if (!messageRunner.isPresent()) {
                 final String errorMessageFormat = "No MessageRunner found to process message (messageID: %s, messageRunnerKey: %s)";
                 final String errorMessage = String.format(errorMessageFormat, NestedMessageConstants.MESSAGE_RUNNER_KEY_ATTRIBUTE_NAME, messageId);
                 log.error(errorMessage);
 
                 return JobRunnerResponse.aborted(errorMessage);
             }
 
             if (!WorkContextAvailable.isWorkContextAvailable()) {
                 final String errorMessageFormat = "No WorkContext available to process message (messageID: %s, messageRunnerKey: %s)";
                 final String errorMessage = String.format(errorMessageFormat, NestedMessageConstants.MESSAGE_RUNNER_KEY_ATTRIBUTE_NAME, messageId);
                 log.error(errorMessage);
 
                 return JobRunnerResponse.aborted(errorMessage);
             }
 
             try {
                 messageRunner.get().processMessage(messageContext);
 
                 if (messageContext.shouldAutoAcknowledgeMessage()) {
                     messageContext.acknowledge();
                 }
             } catch (Exception e) {
                 log.error("Invoking MessageRunner for message failed with an exception (messageID: {}, messageRunnerKey: {})", messageId, messageRunnerKey, e);
             }
 
             return JobRunnerResponse.success();
         }
 
         @Nullable
         JobRunnerResponse validate(JobRunnerRequest jobRunnerRequest) {
             final JobId jobId = jobRunnerRequest.getJobId();
             final Map<String, Serializable> jobParameters = jobRunnerRequest.getJobConfig().getParameters();
             final String payload = (String) jobParameters.get(PAYLOAD_JOB_PARAMETER);
             final Serializable deliveryCount = jobParameters.get(DELIVERY_COUNT_PARAMETER);
 
             if (!jobParameters.containsKey(DELIVERY_COUNT_PARAMETER)) {
                 final String format = "Job missing required parameter: %s. Job will be unscheduled (JobId: %s, JobPayload: %s)";
                 log.error(String.format(format, DELIVERY_COUNT_PARAMETER, jobId, "omitted"));
                 schedulerService.unscheduleJob(jobId);
                 return JobRunnerResponse.aborted(String.format(format, DELIVERY_COUNT_PARAMETER, jobId, payload));
             } else if (!(deliveryCount instanceof Integer)) {
                 final String format = "Job %s parameter of the wrong type. Expecting int, found: %s. Job will be unscheduled (JobId: %s, JobPayload: %s)";
                 log.error(String.format(format, DELIVERY_COUNT_PARAMETER, deliveryCount, jobId, "omitted"));
                 schedulerService.unscheduleJob(jobId);
                 return JobRunnerResponse.aborted(String.format(format, DELIVERY_COUNT_PARAMETER, deliveryCount, jobId, payload));
             }
 
             return null;
         }
 
         @Nullable
         JobRunnerResponse incrementDeliveryCount(JobRunnerRequest jobRunnerRequest) {
             final JobId jobId = jobRunnerRequest.getJobId();
             final Map<String, Serializable> jobParameters = jobRunnerRequest.getJobConfig().getParameters();
             final String payload = (String) jobParameters.get(PAYLOAD_JOB_PARAMETER);
 
             int deliveryCount = (int) jobParameters.get(DELIVERY_COUNT_PARAMETER);
             if (++deliveryCount > maxDeliveryCount) {
                 final String format = "Max delivery count reached for job. Job will be unscheduled (JobId: %s, JobPayload: %s)";
                 log.warn(String.format(format, jobId, "omitted"));
                 schedulerService.unscheduleJob(jobId);
                 return JobRunnerResponse.aborted(String.format(format, jobId, payload));
             } else {
                 final Map<String, Serializable> newJobParameters = new HashMap<>(jobParameters);
                 newJobParameters.put(DELIVERY_COUNT_PARAMETER, deliveryCount);
                 try {
                     schedulerService.scheduleJob(jobId, jobRunnerRequest.getJobConfig()
                             .withSchedule(Schedule.forInterval(deliveryIntervalMillis, getFirstRunTime(jobRunnerRequest.getStartTime())))
                             .withParameters(newJobParameters));
                 } catch (SchedulerServiceException e) {
                     final String format = "Error rescheduling job (JobId: %s, JobPayload: %s)";
                     log.error(String.format(format, jobId, "omitted"), e);
                     return JobRunnerResponse.failed(String.format(format, jobId, payload));
                 }
                 log.debug("JobId: {}, deliveryCount: {}", jobId, deliveryCount);
             }
 
             return null;
         }
     }
 
     private Date getFirstRunTime(Date jobStartTime) {
         return new Date(jobStartTime.getTime() + deliveryIntervalMillis);
     }
 
     /**
      * Shuts down this instance.
      */
     @PreDestroy
     public void shutdown() {
         schedulerService.unregisterJobRunner(AMQ_JOB_RUNNER_KEY);
         shutdown = true;
     }
 
     @Override
     public void addMessage(Message message) {
         if (shutdown) {
             throw new MessageRunnerServiceException("message could not be added after shutdown (runnerKey: " + message.getRunnerKey() + ")");
         }
 
         if (!messageRunnerRegistryHelper.getMessageRunner(message.getRunnerKey()).isPresent()) {
             throw new MessageRunnerNotRegisteredException(message.getRunnerKey());
         }
 
         final Map<String, Serializable> jobParameters = new HashMap<>();
         jobParameters.put(PAYLOAD_JOB_PARAMETER, messageInformationService.toPayload(message));
         jobParameters.put(DELIVERY_COUNT_PARAMETER, 0);
 
         final JobConfig jobConfig = JobConfig.forJobRunnerKey(AMQ_JOB_RUNNER_KEY)
                 .withRunMode(RunMode.RUN_ONCE_PER_CLUSTER)
                 .withSchedule(Schedule.forInterval(deliveryIntervalMillis, null))
                 .withParameters(jobParameters);
 
         try {
             schedulerService.scheduleJobWithGeneratedId(jobConfig);
         } catch (SchedulerServiceException e) {
             log.error("Could not schedule job for message with key {}", message.getRunnerKey().toString(), e);
             throw new MessageRunnerServiceException("Could not add message", e);
         }
     }
 }