/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.activemq.artemis.cli.commands.messages.perf;

import javax.jms.Connection;
import javax.jms.ConnectionFactory;
import javax.jms.Destination;
import javax.jms.JMSException;
import javax.jms.MessageConsumer;
import javax.jms.Session;
import javax.jms.Topic;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.LongAdder;

import org.HdrHistogram.SingleWriterRecorder;

public final class MessageListenerBenchmark implements BenchmarkService {

   private final ConnectionFactory factory;
   private final MicrosTimeProvider timeProvider;
   private final int consumers;
   private final boolean canDelaySetMessageCount;
   private final int connections;
   private final String clientID;
   private final Destination[] destinations;
   private final int sharedSubscription;
   private final boolean durableSubscription;
   private final long messageCount;
   private final boolean transaction;
   private Set<Connection> jmsConnections;
   private boolean started;
   private boolean closed;
   private MessageCountLimiter msgCountLimiter;
   private List<RecordingMessageListener> listeners;
   private AtomicBoolean fatalException;
   private List<Runnable> silentUnsubscribe;

   public static final class MessageCountLimiter {

      private volatile long messageLimit = Long.MAX_VALUE;
      private final LongAdder totalMessagesReceived;

      MessageCountLimiter() {
         totalMessagesReceived = new LongAdder();
      }

      public MessageCountLimiter setMessageLimit(final long messageLimit) {
         this.messageLimit = messageLimit;
         return this;
      }

      public boolean isLimitReached() {
         return totalMessagesReceived.sum() >= messageLimit;
      }

      public void onMessageReceived() {
         totalMessagesReceived.increment();
      }
   }

   public MessageListenerBenchmark(final ConnectionFactory factory,
                                   final MicrosTimeProvider timeProvider,
                                   final int consumers,
                                   final long messageCount,
                                   final int connections,
                                   final String clientID,
                                   final Destination[] destinations,
                                   final boolean transaction,
                                   final int sharedSubscription,
                                   final boolean durableSubscription,
                                   final boolean canDelayMessageCount) {
      this.factory = factory;
      this.timeProvider = timeProvider;
      this.consumers = consumers;
      this.messageCount = messageCount;
      this.connections = connections;
      this.clientID = clientID;
      this.destinations = destinations;
      this.transaction = transaction;
      this.sharedSubscription = sharedSubscription;
      this.durableSubscription = durableSubscription;
      this.started = false;
      this.closed = false;
      this.jmsConnections = new HashSet<>(connections);
      this.canDelaySetMessageCount = canDelayMessageCount;
      this.listeners = null;
      this.fatalException = null;
      this.silentUnsubscribe = null;
   }

   public synchronized RecordingMessageListener[] getListeners() {
      return listeners == null ? null : listeners.toArray(new RecordingMessageListener[listeners.size()]);
   }

   @Override
   public synchronized boolean anyError() {
      if (fatalException == null) {
         return false;
      }
      return fatalException.get();
   }

   @Override
   public synchronized boolean isRunning() {
      if (!started || closed) {
         return false;
      }
      if (fatalException.get()) {
         return false;
      }
      if (msgCountLimiter == null) {
         return true;
      }
      return !msgCountLimiter.isLimitReached();
   }

   public synchronized void setMessageCount(final long messageCount) {
      if (!started || closed) {
         return;
      }
      msgCountLimiter.setMessageLimit(messageCount);
   }

   @Override
   public synchronized MessageListenerBenchmark start() {
      if (started) {
         return this;
      }
      started = true;
      closed = false;
      final AtomicLong consumerId = new AtomicLong(1);
      // setup connection failure listeners
      final AtomicBoolean signalBrokenConnection = new AtomicBoolean(false);
      fatalException = signalBrokenConnection;
      // create connections upfront and register failure listener
      final Connection[] jmsConnections = new Connection[connections];
      for (int i = 0; i < connections; i++) {
         final Connection connection;
         try {
            connection = factory.createConnection();
            if (clientID != null) {
               if (connections > 1) {
                  connection.setClientID(clientID + i);
               } else {
                  connection.setClientID(clientID);
               }
            }
            connection.setExceptionListener(ignore -> {
               signalBrokenConnection.set(true);
            });
            jmsConnections[i] = connection;
         } catch (JMSException e) {
            throw new RuntimeException(e);
         }
      }
      this.jmsConnections.addAll(Arrays.asList(jmsConnections));
      // start connections
      this.jmsConnections.forEach(connection -> {
         try {
            connection.start();
         } catch (JMSException e) {
            throw new RuntimeException(e);
         }
      });
      int connectionSequence = 0;
      final int totalListeners = consumers * destinations.length * Math.max(sharedSubscription, 1);
      this.listeners = new ArrayList<>(totalListeners);
      if (messageCount > 0) {
         msgCountLimiter = new MessageCountLimiter().setMessageLimit(messageCount);
      } else if (canDelaySetMessageCount) {
         msgCountLimiter = new MessageCountLimiter();
      }
      // create consumers per destination
      if (durableSubscription) {
         silentUnsubscribe = new ArrayList<>();
      }
      for (int i = 0; i < destinations.length; i++) {
         final Destination destination = destinations[i];
         if (sharedSubscription == 0) {
            final Queue<RecordingMessageListener> destinationListeners = new ArrayDeque<>(consumers);
            createListeners(destinationListeners, consumerId, destination, consumers);
            listeners.addAll(destinationListeners);
            try {
               for (int consumerIndex = 0; consumerIndex < consumers; consumerIndex++) {
                  final Connection connection = jmsConnections[connectionSequence % connections];
                  connectionSequence++;
                  final Session session = connection.createSession(transaction ? Session.SESSION_TRANSACTED : Session.AUTO_ACKNOWLEDGE);
                  final MessageConsumer consumer;
                  if (durableSubscription) {
                     final Topic topic = (Topic) destination;
                     consumer = session.createDurableConsumer((Topic) destination, topic.getTopicName() + consumerIndex);
                  } else {
                     consumer = session.createConsumer(destination);
                  }
                  consumer.setMessageListener(destinationListeners.remove());
               }
            } catch (JMSException e) {
               throw new RuntimeException(e);
            }
         } else {
            final int listenersPerDestination = sharedSubscription * consumers;
            final Queue<RecordingMessageListener> destinationListeners = new ArrayDeque<>(listenersPerDestination);
            createListeners(destinationListeners, consumerId, destination, listenersPerDestination);
            listeners.addAll(destinationListeners);
            try {
               final String topicName = ((Topic) destination).getTopicName();
               for (int subscriptionIndex = 0; subscriptionIndex < sharedSubscription; subscriptionIndex++) {
                  Connection connection = null;
                  if (clientID != null) {
                     connection = jmsConnections[connectionSequence % connections];
                     assert connection.getClientID() != null;
                     connectionSequence++;
                  }
                  for (int consumerIndex = 0; consumerIndex < consumers; consumerIndex++) {
                     if (clientID == null) {
                        assert connection == null;
                        connection = jmsConnections[connectionSequence % connections];
                        connectionSequence++;
                     }
                     final Session session = connection.createSession(transaction ? Session.SESSION_TRANSACTED : Session.AUTO_ACKNOWLEDGE);
                     final MessageConsumer consumer;
                     if (durableSubscription) {
                        final String subscriptionName = topicName + subscriptionIndex;
                        consumer = session.createSharedDurableConsumer((Topic) destination, subscriptionName);
                        silentUnsubscribe.add(() -> {
                           try {
                              session.unsubscribe(subscriptionName);
                           } catch (JMSException e) {
                              throw new RuntimeException(e);
                           }
                        });
                     } else {
                        consumer = session.createSharedConsumer((Topic) destination, topicName + subscriptionIndex);
                     }
                     consumer.setMessageListener(destinationListeners.remove());
                  }
               }
            } catch (JMSException fatal) {
               throw new RuntimeException(fatal);
            }
         }
      }
      return this;
   }

   private void createListeners(final Collection<? super RecordingMessageListener> listeners,
                                final AtomicLong consumerId,
                                final Destination destination,
                                final int count) {
      for (int c = 0; c < count; c++) {
         listeners.add(new RecordingMessageListener(consumerId.getAndIncrement(), destination, transaction,
                                                    new AtomicLong(0),
                                                    msgCountLimiter == null ? null : msgCountLimiter::onMessageReceived,
                                                    timeProvider, new SingleWriterRecorder(2),
                                                    fatalException));
      }
   }

   @Override
   public synchronized void close() {
      if (!started || closed) {
         return;
      }
      listeners = null;
      started = false;
      closed = true;
      msgCountLimiter = null;
      fatalException = null;
      if (silentUnsubscribe != null) {
         silentUnsubscribe.forEach(Runnable::run);
         silentUnsubscribe = null;
      }
      jmsConnections.forEach(connection -> {
         try {
            connection.close();
         } catch (JMSException ignore) {

         }
      });
      jmsConnections.clear();
   }
}
