/*-
 * -\-\-
 * async-google-pubsub-client
 * --
 * Copyright (C) 2016 - 2017 Spotify AB
 * --
 * Licensed 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.
 * -/-/-
 */

/*
 * Copyright (c) 2011-2015 Spotify AB
 *
 * Licensed 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 com.spotify.google.cloud.pubsub.client;

import com.google.common.util.concurrent.MoreExecutors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Closeable;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import static java.util.Collections.unmodifiableList;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.SECONDS;

/**
 * A tool for publishing larger volumes of messages to Google Pub/Sub. Does concurrent per-topic batching in order to
 * provide good throughput with large volumes of messages across many different topics.
 *
 * <p> Messages are gathered into batches, bounded by specified batch size and maximum latency. The publisher waits up
 * to the specified max latency before sending a batch of messages for a topic. If enough messages to fill a batch are
 * submitted before the max latency deadline, then the batch is sent immediately.
 *
 * <p> This batching strategy trades publish request quota for increased publishing latency as outgoing messages might
 * spend more time waiting in the publisher per-topic queues before getting sent. The rationale for this strategy is to
 * avoid (expensive and empirically observed) excessive numbers of very small batch publish requests during off-peak.
 */
public class Publisher implements Closeable {

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

  /**
   * A listener for monitoring operations performed by the {@link Publisher}.
   */
  public interface Listener {

    /**
     * Called when a new {@link Publisher} is instantiated.
     *
     * @param publisher The {@link Publisher}
     */
    void publisherCreated(Publisher publisher);

    /**
     * Called when a {@link Publisher} is closed.
     *
     * @param publisher The {@link Publisher}
     */
    void publisherClosed(Publisher publisher);

    /**
     * Called when a {@link Publisher} receieves a new message for publication.
     *
     * @param publisher The {@link Publisher}
     * @param topic     The message topic.
     * @param message   The message.
     * @param future    The future result.
     */
    void publishingMessage(Publisher publisher, String topic, Message message, CompletableFuture<String> future);

    /**
     * Called when a {@link Publisher} is sending a batch of messages to Google Cloud Pub/Sub.
     *
     * @param publisher The {@link Publisher}
     * @param topic     The topic of the message batch.
     * @param batch     The batch of messages being sent.
     * @param future    The future result of the entire batch.
     */
    @Deprecated
    void sendingBatch(Publisher publisher, String topic, List<Message> batch, CompletableFuture<List<String>> future);

    /**
     * Called when a {@link Publisher} is sending a batch of messages to Google Cloud Pub/Sub.
     *
     * @param publisher The {@link Publisher}
     * @param topic     The topic of the message batch.
     * @param batch     The batch of messages being sent.
     * @param future    The future result of the entire batch.
     */
    default void sendingBatch(Publisher publisher, String topic, List<Message> batch,
                              PubsubFuture<List<String>> future) {
      sendingBatch(publisher, topic, batch, (CompletableFuture<List<String>>) future);
    }

    /**
     * Called when a topic is enqueued as pending for future batch sending due to the publisher hitting the concurrency
     * limit.
     *
     * @param publisher   The {@link Publisher}
     * @param topic       The topic.
     * @param outstanding The current number of outstanding batch requests to Google Cloud Pub/Sub.
     * @param concurrency The configured concurrency limit.
     */
    void topicPending(Publisher publisher, String topic, int outstanding, int concurrency);
  }

  private final Pubsub pubsub;
  private final String project;
  private final int queueSize;
  private final int batchSize;
  private final int concurrency;
  private final long maxLatencyMs;
  private final Listener listener;

  private final AtomicInteger outstanding = new AtomicInteger();
  private final ConcurrentLinkedQueue<TopicQueue> pendingTopics = new ConcurrentLinkedQueue<>();
  private final ConcurrentMap<String, TopicQueue> topics = new ConcurrentHashMap<>();
  private final CompletableFuture<Void> closeFuture = new CompletableFuture<>();
  private final ScheduledExecutorService scheduler =
      MoreExecutors.getExitingScheduledExecutorService(new ScheduledThreadPoolExecutor(1));
  private final AtomicInteger schedulerQueueSize = new AtomicInteger();

  private Publisher(final Builder builder) {
    this.pubsub = Objects.requireNonNull(builder.pubsub, "pubsub");
    this.project = Objects.requireNonNull(builder.project, "project");
    this.concurrency = builder.concurrency;
    this.batchSize = builder.batchSize;
    this.queueSize = Optional.ofNullable(builder.queueSize).orElseGet(() -> batchSize * 10);
    this.maxLatencyMs = builder.maxLatencyMs;
    this.listener = builder.listener == null ? new ListenerAdapter() : builder.listener;
    listener.publisherCreated(this);
  }

  /**
   * Publish a message on a specific topic.
   *
   * @param topic   The topic name to publish on. Note that this is the short name, not the fully qualified name
   *                including project. The project to publish on is configured using the {@link Builder}.
   * @param message The message to publish.
   * @return A future that is fulfilled with the resulting Google Pub/Sub message ID when the message has been
   * successfully published.
   */
  public CompletableFuture<String> publish(final String topic, final Message message) {
    final TopicQueue queue = topics.computeIfAbsent(topic, TopicQueue::new);
    final CompletableFuture<String> future = queue.send(message);
    listener.publishingMessage(this, topic, message, future);
    return future;
  }

  /**
   * Close this {@link Publisher}. This will also close the underlying {@link Pubsub} client.
   */
  @Override
  public void close() {
    // TODO (dano): fail outstanding futures
    scheduler.shutdownNow();
    try {
      scheduler.awaitTermination(30, SECONDS);
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
    }
    closeFuture.complete(null);
    listener.publisherClosed(Publisher.this);
  }

  /**
   * Get a future that is completed when this {@link Publisher} is closed.
   */
  public CompletableFuture<Void> closeFuture() {
    return closeFuture.thenApply(ignore -> null);
  }

  /**
   * Get the current number of outstanding (batch) requests to Google Pub/Sub.
   */
  public int outstandingRequests() {
    return outstanding.get();
  }

  /**
   * Get the total number of scheduled requests.
   */
  public int scheduledQueueSize() {
    return schedulerQueueSize.get();
  }

  /**
   * Get the current number of topics that are pending batch sending to Google Pub/Sub.
   */
  public int pendingTopics() {
    return pendingTopics.size();
  }

  /**
   * Get the Google Cloud project this {@link Publisher} is publishing to.
   */
  public String project() {
    return project;
  }

  /**
   * Get the concurrent Google Pub/Sub request limit.
   */
  public int concurrency() {
    return concurrency;
  }

  /**
   * Get the current queue size for the given topic name. Returns 0 if the topic does not exist.
   *
   * @param topic  the topic name
   */
  public int topicQueueSize(final String topic) {
    final TopicQueue topicQueue = this.topics.get(topic);
    return topicQueue == null ? 0 : topicQueue.queue.size();
  }

  /**
   * Get the per-topic queue size.
   */
  public int queueSize() {
    return queueSize;
  }

  /**
   * Get the max Google Pub/Sub request batch size.
   */
  public int batchSize() {
    return batchSize;
  }

  /**
   * The per-topic queue of messages.
   */
  private class TopicQueue {

    private final AtomicInteger size = new AtomicInteger();
    private final ConcurrentLinkedQueue<QueuedMessage> queue = new ConcurrentLinkedQueue<>();
    private final String topic;

    private volatile boolean pending;
    private final AtomicBoolean scheduled = new AtomicBoolean();

    private TopicQueue(final String topic) {
      this.topic = topic;
    }

    /**
     * Enqueue a message for sending on this topic queue.
     */
    private CompletableFuture<String> send(final Message message) {
      final CompletableFuture<String> future = new CompletableFuture<>();

      // Enforce queue size limit
      int currentSize;
      int newSize;
      do {
        currentSize = size.get();
        newSize = currentSize + 1;
        if (newSize > queueSize) {
          future.completeExceptionally(new QueueFullException());
          return future;
        }
      } while (!size.compareAndSet(currentSize, newSize));

      // Enqueue outgoing message
      queue.add(new QueuedMessage(message, future));

      // Schedule future batch sending
      scheduleSend(newSize);

      return future;
    }

    /**
     * Schedule this topic for future enqueuing for batch sending. If the batch size has been reached, enqueue for
     * sending immediately.
     *
     * @param queueSize The current number of enqueued messages in this topic.
     */
    private void scheduleSend(final int queueSize) {

      // Bail if this topic is already enqueued for sending.
      if (pending) {
        return;
      }

      // Reached the batch size? Enqueue topic for sending immediately.
      if (queueSize >= batchSize) {
        enqueueSend();
        return;
      }

      // Schedule this topic for later enqueuing, allowing more messages to gather into a larger batch.
      if (scheduled.compareAndSet(false, true)) {
        try {
          scheduler.schedule(this::scheduledEnqueueSend, maxLatencyMs, MILLISECONDS);
          schedulerQueueSize.incrementAndGet();
        } catch (RejectedExecutionException ignore) {
          // Race with a call to close(). Ignore.
        }
      }
    }

    /**
     * Decrements the scheduled queue counter and enqueues the request.
     */
    private void scheduledEnqueueSend() {
      schedulerQueueSize.decrementAndGet();
      // Clear the scheduled flag before enqueuing or sending.
      scheduled.set(false);
      enqueueSendWithErrorLogging();
    }

    /**
     * A wrapper around enqueueSend which catches and logs any exceptions that are thrown. This is
     * called by the executor, which will silently swallow exceptions if we don't handle them here.
     */
    private void enqueueSendWithErrorLogging() {
      try {
        enqueueSend();
      } catch (Exception e) {
        log.error("Error while enqueueing or sending messages on background thread", e);
      }
    }

    /**
     * Enqueue this topic for batch sending. If the request concurrency level is below the limit, send immediately.
     */
    private void enqueueSend() {

      final int currentOutstanding = outstanding.get();

      // Below outstanding limit so we can send immediately?
      if (currentOutstanding < concurrency) {
        sendBatch();
        return;
      }

      // Enqueue as pending for sending by the earliest available concurrent request slot
      pending = true;
      pendingTopics.offer(this);

      // Tell the listener that a topic became pending for sending as early as possible.
      listener.topicPending(Publisher.this, topic, currentOutstanding, concurrency);

      // Attempt to send pending to guard against losing a race while enqueuing this topic as pending.
      sendPending();
    }

    /**
     * Send a batch of messages.
     */
    private int sendBatch() {
      final List<Message> batch = new ArrayList<>();
      final List<CompletableFuture<String>> futures = new ArrayList<>();

      // Drain queue up to batch size
      while (batch.size() < batchSize) {
        final QueuedMessage message = queue.poll();
        if (message == null) {
          break;
        }
        batch.add(message.message);
        futures.add(message.future);
      }

      // Was there anything to send?
      if (batch.size() == 0) {
        return 0;
      }

      // Decrement the queue size counter
      size.updateAndGet(i -> i - batch.size());

      // Send the batch request and increment the outstanding request counter
      outstanding.incrementAndGet();
      final PubsubFuture<List<String>> batchFuture = pubsub.publish(project, topic, batch);
      listener.sendingBatch(Publisher.this, topic, unmodifiableList(batch), batchFuture);
      batchFuture.whenComplete(
          (List<String> messageIds, Throwable ex) -> {

            // Decrement the outstanding request counter
            outstanding.decrementAndGet();

            // Fail all futures if the batch request failed
            if (ex != null) {
              futures.forEach(f -> f.completeExceptionally(ex));
              return;
            }

            // Verify that the number of message id's and messages match up
            if (futures.size() != messageIds.size()) {
              futures.forEach(f -> f.completeExceptionally(
                  new PubsubException(
                      "message id count mismatch: " +
                      futures.size() + " != " + messageIds.size())));
            }

            // Complete each future with the appropriate message id
            for (int i = 0; i < futures.size(); i++) {
              final String messageId = messageIds.get(i);
              final CompletableFuture<String> future = futures.get(i);
              future.complete(messageId);
            }
          })

          // When batch is complete, process pending topics.
          .whenComplete((v, t) -> sendPending());

      return batch.size();
    }
  }

  /**
   * Send any pending topics.
   */
  private void sendPending() {
    while (outstanding.get() < concurrency) {
      final TopicQueue queue = pendingTopics.poll();
      if (queue == null) {
        return;
      }

      queue.pending = false;
      final int sent = queue.sendBatch();

      // Did we send a whole batch? Then there might be more messages in the queue. Mark as pending again.
      if (sent == batchSize) {
        queue.pending = true;
        pendingTopics.offer(queue);
      }
    }
  }

  /**
   * An outgoing message with the future that should be completed when the message has been published.
   */
  private static class QueuedMessage {

    private final Message message;
    private final CompletableFuture<String> future;

    public QueuedMessage(final Message message, final CompletableFuture<String> future) {
      this.message = message;
      this.future = future;
    }
  }

  /**
   * Create a builder that can be used to build a {@link Publisher}.
   */
  public static Builder builder() {
    return new Builder();
  }

  /**
   * A builder that can be used to build a {@link Publisher}.
   */
  public static class Builder {

    private Pubsub pubsub;
    private String project;
    private Integer queueSize;
    private int batchSize = 1000;
    private int concurrency = 64;
    private Listener listener;
    private long maxLatencyMs = 100;

    /**
     * Set the {@link Pubsub} client to use.
     *
     * <p>Note: The client should be configured to at least allow as many connections as the concurrency level of this
     * {@link Publisher}.</p>
     */
    public Builder pubsub(final Pubsub pubsub) {
      this.pubsub = pubsub;
      return this;
    }

    /**
     * Set the Google Cloud project to publish to.
     */
    public Builder project(final String project) {
      this.project = project;
      return this;
    }

    /**
     * Set the maximum batch size. Default is {@code 1000}, which is also the maximum Google Cloud Pub/Sub batch size.
     */
    public Builder batchSize(final int batchSize) {
      this.batchSize = batchSize;
      return this;
    }

    /**
     * Set the per-topic queue size. Default is {@code batchSize * concurrency * 10}.
     */
    public Builder queueSize(final Integer queueSize) {
      this.queueSize = queueSize;
      return this;
    }

    /**
     * Set the Google Cloud Pub/Sub request concurrency level. Default is {@code 64}.
     */
    public Builder concurrency(final int concurrency) {
      this.concurrency = concurrency;
      return this;
    }

    /**
     * Set the maximum latency in millis before sending an incomplete Google Cloud Pub/Sub publish batch request.
     * Default is {@code 100 ms}.
     */
    public Builder maxLatencyMs(final long maxLatencyMs) {
      this.maxLatencyMs = maxLatencyMs;
      return this;
    }

    /**
     * Set a {@link Listener} for monitoring operations performed by the {@link Publisher}.
     */
    public Builder listener(final Listener listener) {
      this.listener = listener;
      return this;
    }

    /**
     * Build a {@link Publisher}.
     */
    public Publisher build() {
      return new Publisher(this);
    }
  }

  public static class ListenerAdapter implements Listener {

    @Override
    public void publisherCreated(final Publisher publisher) {

    }

    @Override
    public void publisherClosed(final Publisher publisher) {

    }

    @Override
    public void publishingMessage(final Publisher publisher, final String topic, final Message message,
                                  final CompletableFuture<String> future) {
    }

    @Override
    public void sendingBatch(final Publisher publisher, final String topic, final List<Message> batch,
                             final PubsubFuture<List<String>> future) {
    }

    @Override
    public void sendingBatch(final Publisher publisher, final String topic, final List<Message> batch,
                             final CompletableFuture<List<String>> future) {

    }

    @Override
    public void topicPending(final Publisher publisher, final String topic, final int outstanding,
                             final int concurrency) {

    }
  }
}
