package com.mozilla.telemetry.ingestion.sink.io;

import com.google.api.core.ApiFuture;
import com.google.api.core.ApiFutureCallback;
import com.google.api.core.ApiFutures;
import com.google.cloud.pubsub.v1.Publisher;
import com.google.cloud.pubsub.v1.Subscriber;
import com.google.common.annotations.VisibleForTesting;
import com.google.pubsub.v1.ProjectSubscriptionName;
import com.google.pubsub.v1.ProjectTopicName;
import com.google.pubsub.v1.PubsubMessage;
import com.mozilla.telemetry.ingestion.sink.transform.PubsubMessageToTemplatedString;
import com.mozilla.telemetry.ingestion.sink.util.BatchException;
import java.io.IOException;
import java.io.UncheckedIOException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executor;
import java.util.function.Function;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class Pubsub {

  private Pubsub() {
  }

  public static class Read {

    private static final Logger LOG = LoggerFactory.getLogger(Read.class);

    @VisibleForTesting
    public Subscriber subscriber;

    /** Constructor. */
    public <T> Read(String subscriptionName, Function<PubsubMessage, CompletableFuture<T>> output,
        Function<Subscriber.Builder, Subscriber.Builder> config,
        Function<PubsubMessage, PubsubMessage> decompress) {
      ProjectSubscriptionName subscription = ProjectSubscriptionName.parse(subscriptionName);
      subscriber = config.apply(Subscriber.newBuilder(subscription,
          // Synchronous CompletableFuture methods are executed by the thread that completes the
          // future, or the current thread if the future is already complete. Use that here to
          // minimize memory usage by doing as much work as immediately possible.
          (message, consumer) -> CompletableFuture.completedFuture(message).thenApply(decompress)
              .thenCompose(output).whenComplete((result, exception) -> {
                if (exception == null) {
                  consumer.ack();
                } else {
                  // exception is always a CompletionException caused by another exception
                  if (exception.getCause() instanceof BatchException) {
                    // only log batch exception once
                    ((BatchException) exception.getCause()).handle((batchExc) -> LOG.error(
                        String.format("failed to deliver %d messages", batchExc.size),
                        batchExc.getCause()));
                  } else {
                    // log exception specific to this message
                    LOG.error("failed to deliver message", exception.getCause());
                  }
                  consumer.nack();
                }
              })))
          .build();
    }

    /** Run the subscriber until terminated. */
    public void run() {
      try {
        subscriber.startAsync();
        subscriber.awaitTerminated();
      } finally {
        subscriber.stopAsync();
      }
    }
  }

  public static class Write implements Function<PubsubMessage, CompletableFuture<String>> {

    private final Executor executor;
    private final Function<Publisher.Builder, Publisher.Builder> config;
    private final Function<PubsubMessage, PubsubMessage> compress;
    private final PubsubMessageToTemplatedString topicTemplate;
    private final ConcurrentMap<String, Publisher> publishers = new ConcurrentHashMap<>();

    /** Constructor. */
    public Write(String topicTemplate, Executor executor,
        Function<Publisher.Builder, Publisher.Builder> config,
        Function<PubsubMessage, PubsubMessage> compress) {
      this.executor = executor;
      this.topicTemplate = PubsubMessageToTemplatedString.of(topicTemplate);
      this.config = config;
      this.compress = compress;
    }

    private Publisher getPublisher(PubsubMessage message) {
      return publishers.compute(topicTemplate.apply(message), (topic, publisher) -> {
        if (publisher == null) {
          try {
            return config.apply(Publisher.newBuilder(ProjectTopicName.parse(topic))).build();
          } catch (IOException e) {
            throw new UncheckedIOException(e);
          }
        }
        return publisher;
      });
    }

    @Override
    public CompletableFuture<String> apply(PubsubMessage message) {
      final PubsubMessage compressed = compress.apply(message);
      final ApiFuture<String> future = getPublisher(message).publish(compressed);
      final CompletableFuture<String> result = new CompletableFuture<>();
      ApiFutures.addCallback(future, new ApiFutureCallback<String>() {

        @Override
        public void onFailure(Throwable throwable) {
          result.completeExceptionally(throwable);
        }

        @Override
        public void onSuccess(String messageId) {
          result.complete(messageId);
        }
      }, executor);
      return result;
    }

    public CompletableFuture<Void> withoutResult(PubsubMessage message) {
      return apply(message).thenAccept(result -> {
      });
    }
  }
}