package com.atlassian.messagequeue.internal.sqs;
   
   import com.amazonaws.services.sqs.AmazonSQS;
   import com.amazonaws.services.sqs.model.InvalidMessageContentsException;
   import com.amazonaws.services.sqs.model.SendMessageRequest;
   import com.amazonaws.services.sqs.model.SendMessageResult;
   import com.amazonaws.services.sqs.model.UnsupportedOperationException;
   import com.atlassian.messagequeue.Message;
   import com.atlassian.messagequeue.MessageInformationService;
  import com.atlassian.messagequeue.MessageRunnerNotRegisteredException;
  import com.atlassian.messagequeue.MessageRunnerService;
  import com.atlassian.messagequeue.MessageRunnerServiceException;
  import com.atlassian.messagequeue.TenantDataIdSupplier;
  import com.atlassian.messagequeue.internal.core.DefaultThreadFactory;
  import com.atlassian.messagequeue.internal.core.MessageRunnerRegistryHelper;
  import com.atlassian.messagequeue.internal.core.MessageValidatorRegistryHelper;
  import com.atlassian.messagequeue.internal.core.NestedMessageConsumer;
  import com.atlassian.messagequeue.internal.core.NestedMessageSerializer;
  import com.atlassian.tenant.api.TenantContextProvider;
  import com.atlassian.tenant.impl.TenantIdSetter;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.slf4j.MDC;
  
  import javax.annotation.Nonnull;
  import javax.annotation.PreDestroy;
  import java.util.Collections;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.Set;
  import java.util.concurrent.ExecutorService;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.ThreadPoolExecutor;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicBoolean;
  
  import static com.atlassian.messagequeue.internal.core.NestedMessageConsumer.MDC_MESSAGE_RUNNER_KEY;
  import static java.util.Objects.requireNonNull;
  
  /**
   * SQS message runner service that produces messages onto an SQS queue and manages a pool of threads that consume from said queue.
   * <p>
   * To properly initialise an instance of this service, a client must call {@link #initialiseMessageConsumers()} 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 SQSMessageRunnerService implements MessageRunnerService {
      static final String SENT_TIMESTAMP = "SentTimestamp";
      static final String APPROXIMATE_RECEIVE_COUNT = "ApproximateReceiveCount";
      static final long AWAIT_TERMINATION_TIMEOUT_SECONDS = Long.getLong("amq.sqs.await.termination.timeout", TimeUnit.SECONDS.convert(1, TimeUnit.HOURS));
      static final String MDC_MESSAGE_ID = "amq.messageId";
  
      private static final Logger log = LoggerFactory.getLogger(SQSMessageRunnerService.class);
  
      private static final int DEFAULT_RECEIVE_WAIT_TIME_SECONDS = Integer.getInteger("amq.sqs.receive.wait.time", 20);
  
      private final AmazonSQS amazonSQSClient;
      private final MessageRunnerRegistryHelper messageRunnerRegistryHelper;
      private final AtomicBoolean shuttingDown = new AtomicBoolean(false);
      private final int receiveWaitTimeSeconds;
      private final Map<String, ExecutorService> queueNameToConsumerThreadPools = new HashMap<>();
      private final AtomicBoolean initialised = new AtomicBoolean(false);
      private final MessageInformationService messageInformationService;
      private final NestedMessageSerializer nestedMessageSerializer;
      private final NestedMessageConsumer nestedMessageConsumer;
      private final SQSMessageVisibilityTimeoutManager sqsMessageVisibilityTimeoutManager;
      private final SQSMessageRunnerKeyToProducerMapper messageRunnerKeyToProducerMapper;
      private final Set<SQSConsumerQueueConfig> consumerQueueConfigs;
  
      private SQSMessageRunnerService(Builder builder) {
          if (builder.messageValidatorRegistryHelper == null) {
              builder.messageValidatorRegistryHelper = Collections::emptyList;
          }
  
          if (builder.tenantDataIdSupplier == null) {
              builder.tenantDataIdSupplier = () -> "";
          }
  
          if (builder.receiveWaitTimeSeconds == null) {
              builder.receiveWaitTimeSeconds = DEFAULT_RECEIVE_WAIT_TIME_SECONDS;
          }
  
          if (builder.sqsMessageVisibilityTimeoutManager == null) {
              builder.sqsMessageVisibilityTimeoutManager = new SQSMessageVisibilityTimeoutManager(
                      requireNonNull(builder.amazonSQSClient));
          }
  
          this.amazonSQSClient = requireNonNull(builder.amazonSQSClient);
          this.messageRunnerRegistryHelper = requireNonNull(builder.messageRunnerRegistryHelper);
          this.receiveWaitTimeSeconds = builder.receiveWaitTimeSeconds;
          this.nestedMessageSerializer = requireNonNull(builder.nestedMessageSerializer);
          this.nestedMessageConsumer = new NestedMessageConsumer(messageRunnerRegistryHelper, requireNonNull(builder.tenantIdSetter), requireNonNull(builder.messageValidatorRegistryHelper));
          this.messageRunnerKeyToProducerMapper = requireNonNull(builder.messageRunnerKeyToProducerMapper);
          this.consumerQueueConfigs = requireNonNull(builder.consumerQueueConfigs);
  
         // Create separate thread pools for configured queues
         consumerQueueConfigs.forEach(queueConfig -> {
             // Create ThreadPools for the queue
             final ThreadPoolExecutor queueExecutor = new ThreadPoolExecutor(queueConfig.getCorePoolSize(), queueConfig.getMaxPoolSize(), 0L, TimeUnit.MILLISECONDS,
                     new LinkedBlockingQueue<>(),
                     new DefaultThreadFactory("sqs-consumer-thread-" + queueConfig.getQueueName() + "-%d", (t, throwable) -> {
                         log.warn("SQS consumer thread '{}' died due to an exception.", t.getName(), throwable);
                     }));
             queueNameToConsumerThreadPools.put(queueConfig.getQueueName(), queueExecutor);
         });
 
         if (builder.messageInformationService == null) {
             builder.messageInformationService = new SQSMessageInformationService(messageRunnerKeyToProducerMapper,
                     requireNonNull(builder.tenantContextProvider, "Either provide a TenantContextProvider, or a MessageInformationService"),
                     requireNonNull(builder.nestedMessageSerializer, "Either provide a NestedMessageSerializer, or a MessageInformationService"),
                     requireNonNull(builder.tenantDataIdSupplier));
         }
         this.messageInformationService = requireNonNull(builder.messageInformationService);
 
         this.sqsMessageVisibilityTimeoutManager = requireNonNull(builder.sqsMessageVisibilityTimeoutManager, "sqsMessageVisibilityTimeoutManager");
     }
 
     public synchronized void initialiseMessageConsumers() {
         if (initialised.get()) {
             log.info("SQS consumers have already been initialised. Skipping it.");
             return;
         }
 
         log.info("Initialising SQS message consumers");
 
         consumerQueueConfigs.forEach(queueConfig -> {
             log.info("Constructing queue consumer thread pool (queueName: {}, concurrentConsumers:  {}, queueUrl: {}, receiveWaitTimeSeconds: {})",
                     queueConfig.getQueueName(), queueConfig.getCorePoolSize(), queueConfig.getQueueUrl(), receiveWaitTimeSeconds);
 
             // Spawn message consumer threads
             final ExecutorService executorService = queueNameToConsumerThreadPools.get(queueConfig.getQueueName());
             for (int i = 0; i < queueConfig.getCorePoolSize(); i++) {
                 executorService.execute(new SQSMessageConsumer(queueConfig, receiveWaitTimeSeconds, shuttingDown,
                         amazonSQSClient, nestedMessageSerializer, nestedMessageConsumer, sqsMessageVisibilityTimeoutManager));
             }
         });
 
         initialised.compareAndSet(false, true);
     }
 
     @PreDestroy
     public void shutdown() {
         log.info("Shutting down {}", this.getClass().getSimpleName());
 
         shuttingDown.compareAndSet(false, true);
 
         // Tell all executors to shutdown
         queueNameToConsumerThreadPools.forEach((queueName, executor) -> executor.shutdown());
         // Wait for all executors to shutdown
         queueNameToConsumerThreadPools.forEach((queueName, executor) -> {
             try {
                 if (!executor.awaitTermination(AWAIT_TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
                     log.warn("Pool for queue {} did not terminate in {} seconds", queueName, AWAIT_TERMINATION_TIMEOUT_SECONDS);
                 }
             } catch (InterruptedException ie) {
                 Thread.currentThread().interrupt();
             }
         });
 
         sqsMessageVisibilityTimeoutManager.shutdown();
     }
 
     @Override
     public void addMessage(@Nonnull Message message) {
         if (!messageRunnerRegistryHelper.getMessageRunner(message.getRunnerKey()).isPresent()) {
             throw new MessageRunnerNotRegisteredException(message.getRunnerKey());
         }
 
         final SendMessageResult sendMessageResult;
         final String queueUrl = messageInformationService.getQueueUrl(message.getRunnerKey());
         try {
             sendMessageResult = amazonSQSClient.sendMessage(new SendMessageRequest(queueUrl,
                     messageInformationService.toPayload(message)));
         } catch (InvalidMessageContentsException e) {
             throw new MessageRunnerServiceException(e);
         } catch (UnsupportedOperationException e) {
             throw new java.lang.UnsupportedOperationException(e);
         }
 
         if (log.isInfoEnabled()) {
             MDC.put(MDC_MESSAGE_ID, sendMessageResult.getMessageId());
             MDC.put(MDC_MESSAGE_RUNNER_KEY, message.getRunnerKey().toString());
             try {
                 log.info("Message produced to {} (messageId: {}, messageRunnerKey: {})",
                         queueUrl, sendMessageResult.getMessageId(), message.getRunnerKey().toString());
             } finally {
                 MDC.remove(MDC_MESSAGE_ID);
                 MDC.remove(MDC_MESSAGE_RUNNER_KEY);
             }
         }
     }
 
     /**
      * Create a new SQSMessageRunnerService builder
      *
      * @param messageRunnerKeyToProducerMapper mapping from {@link com.atlassian.messagequeue.MessageRunnerKey}s to {@link SQSProducerQueueConfig}s
      * @param consumerQueueConfigs             all {@link SQSConsumerQueueConfig}s should be listened to.
      * @since 2.1.0
      */
     public static Builder newBuilder(SQSMessageRunnerKeyToProducerMapper messageRunnerKeyToProducerMapper, Set<SQSConsumerQueueConfig> consumerQueueConfigs) {
         return new Builder(messageRunnerKeyToProducerMapper, consumerQueueConfigs);
     }
 
     /**
      * A builder to help make the SQSMessageRunnerService.
      * <p>
      * The {@link TenantDataIdSupplier} and {@link MessageValidatorRegistryHelper} are optional.
      * You may choose to provide a {@link MessageInformationService} instead of a {@link TenantContextProvider}, and a {@link NestedMessageSerializer}.
      * You may optionally provide a {@link SQSMessageVisibilityTimeoutManager}
      */
     public static class Builder {
         private AmazonSQS amazonSQSClient;
         private MessageRunnerRegistryHelper messageRunnerRegistryHelper;
         private TenantIdSetter tenantIdSetter;
         private MessageInformationService messageInformationService;
         private NestedMessageSerializer nestedMessageSerializer;
         private MessageValidatorRegistryHelper messageValidatorRegistryHelper;
         private Integer receiveWaitTimeSeconds;
         private SQSMessageVisibilityTimeoutManager sqsMessageVisibilityTimeoutManager;
         private TenantContextProvider<?> tenantContextProvider;
         private TenantDataIdSupplier tenantDataIdSupplier;
         private SQSMessageRunnerKeyToProducerMapper messageRunnerKeyToProducerMapper;
         private Set<SQSConsumerQueueConfig> consumerQueueConfigs;
 
         Builder(SQSMessageRunnerKeyToProducerMapper messageRunnerKeyToProducerMapper, Set<SQSConsumerQueueConfig> consumerQueueConfigs) {
             this.messageRunnerKeyToProducerMapper = requireNonNull(messageRunnerKeyToProducerMapper);
             this.consumerQueueConfigs = new HashSet<>(requireNonNull(consumerQueueConfigs));
         }
 
         public Builder withAmazonSQSClient(AmazonSQS amazonSQSClient) {
             this.amazonSQSClient = amazonSQSClient;
             return this;
         }
 
         public Builder withMessageRunnerRegistryHelper(MessageRunnerRegistryHelper messageRunnerRegistryHelper) {
             this.messageRunnerRegistryHelper = messageRunnerRegistryHelper;
             return this;
         }
 
         public Builder withTenantIdSetter(TenantIdSetter tenantIdSetter) {
             this.tenantIdSetter = tenantIdSetter;
             return this;
         }
 
         public Builder withMessageInformationService(MessageInformationService messageInformationService) {
             this.messageInformationService = messageInformationService;
             return this;
         }
 
         public Builder withNestedMessageSerializer(NestedMessageSerializer nestedMessageSerializer) {
             this.nestedMessageSerializer = nestedMessageSerializer;
             return this;
         }
 
         public Builder withMessageValidatorRegistryHelper(MessageValidatorRegistryHelper messageValidatorRegistryHelper) {
             this.messageValidatorRegistryHelper = messageValidatorRegistryHelper;
             return this;
         }
 
         /* For testing only */
         Builder withReceiveWaitTimeSeconds(int receiveWaitTimeSeconds) {
             this.receiveWaitTimeSeconds = receiveWaitTimeSeconds;
             return this;
         }
 
         public Builder withSqsMessageVisibilityTimeoutManager(SQSMessageVisibilityTimeoutManager sqsMessageVisibilityTimeoutManager) {
             this.sqsMessageVisibilityTimeoutManager = sqsMessageVisibilityTimeoutManager;
             return this;
         }
 
         public Builder withTenantContextProvider(TenantContextProvider<?> tenantContextProvider) {
             this.tenantContextProvider = tenantContextProvider;
             return this;
         }
 
         public Builder withTenantDataIdSupplier(TenantDataIdSupplier tenantDataIdSupplier) {
             this.tenantDataIdSupplier = tenantDataIdSupplier;
             return this;
         }
 
         public SQSMessageRunnerService build() {
             return new SQSMessageRunnerService(this);
         }
     }
 
 
 }