/*
 * Copyright 2013-2021 Signal Messenger, LLC
 * SPDX-License-Identifier: AGPL-3.0-only
 */

package org.whispersystems.textsecuregcm.storage;

import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;

import com.google.common.annotations.VisibleForTesting;
import io.dropwizard.lifecycle.Managed;
import io.lettuce.core.cluster.models.partitions.RedisClusterNode;
import io.micrometer.core.instrument.Counter;
import io.micrometer.core.instrument.DistributionSummary;
import io.micrometer.core.instrument.LongTaskTimer;
import io.micrometer.core.instrument.Metrics;
import io.micrometer.core.instrument.Tag;
import io.micrometer.core.instrument.Tags;
import io.micrometer.core.instrument.Timer;
import java.time.Clock;
import java.time.Duration;
import java.time.Instant;
import java.util.Locale;
import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicLong;
import javax.annotation.Nullable;
import org.apache.commons.lang3.tuple.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
import org.whispersystems.textsecuregcm.entities.MessageProtos;
import org.whispersystems.textsecuregcm.identity.IdentityType;
import org.whispersystems.textsecuregcm.metrics.DevicePlatformUtil;
import org.whispersystems.textsecuregcm.util.Util;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Scheduler;
import reactor.util.function.Tuples;
import reactor.util.retry.Retry;
import reactor.util.retry.RetryBackoffSpec;
import software.amazon.awssdk.services.dynamodb.model.ItemCollectionSizeLimitExceededException;

public class MessagePersister implements Managed {

  private final MessagesCache messagesCache;
  private final MessagesManager messagesManager;
  private final AccountsManager accountsManager;
  private final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager;
  private final Scheduler persistQueueScheduler;
  private final Clock clock;

  private final Duration persistDelay;
  private final int maxConcurrency;
  private final int scanCount;

  private final RetryBackoffSpec retryBackoffSpec;

  private final String persisterId = UUID.randomUUID().toString();

  private volatile boolean running;

  @Nullable
  private Thread workerThread;

  private static final String INSPECTED_QUEUE_COUNTER_NAME = name(MessagePersister.class, "inspectedQueue");

  private static final String OVERSIZED_QUEUE_COUNTER_NAME = name(MessagePersister.class, "persistQueueOversized");
  private static final String PERSISTED_MESSAGE_COUNTER_NAME = name(MessagePersister.class, "persistMessage");
  private static final String PERSISTED_BYTES_COUNTER_NAME = name(MessagePersister.class, "persistBytes");

  private static final Timer PERSIST_QUEUE_TIMER = Metrics.timer(name(MessagePersister.class, "persistQueue"));
  private static final Counter TRIMMED_MESSAGE_COUNTER = Metrics.counter(name(MessagePersister.class, "trimmedMessage"));
  private static final Counter TRIMMED_MESSAGE_BYTES_COUNTER = Metrics.counter(name(MessagePersister.class, "trimmedMessageBytes"));

  private static final Counter PERSIST_NODE_COUNTER = Metrics.counter(name(MessagePersister.class, "persistNodeCount"));
  private static final LongTaskTimer PERSIST_NODE_TIMER = Metrics.more().longTaskTimer(name(MessagePersister.class, "persistNode"));

  private static final String QUEUE_SIZE_DISTRIBUTION_SUMMARY_NAME = name(MessagePersister.class, "queueSize");

  static final int QUEUE_BATCH_LIMIT = 100;
  static final int MESSAGE_BATCH_LIMIT = 100;

  private static final DistributionSummary QUEUE_COUNT_DISTRIBUTION_SUMMARY =
      DistributionSummary.builder(name(MessagePersister.class, "queueCount"))
          .register(Metrics.globalRegistry);

  private static final long EXCEPTION_PAUSE_MILLIS = Duration.ofSeconds(3).toMillis();

  private static final Logger logger = LoggerFactory.getLogger(MessagePersister.class);

  public MessagePersister(final MessagesCache messagesCache,
      final MessagesManager messagesManager,
      final AccountsManager accountsManager,
      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
      final Scheduler persistQueueScheduler,
      final Clock clock,
      final Duration persistDelay,
      final int maxConcurrency,
      final int scanCount) {

    this(messagesCache,
        messagesManager,
        accountsManager,
        dynamicConfigurationManager,
        persistQueueScheduler,
        clock,
        persistDelay,
        maxConcurrency,
        scanCount,
        Retry.backoff(3, Duration.ofSeconds(1)));
  }

  @VisibleForTesting
  MessagePersister(final MessagesCache messagesCache,
      final MessagesManager messagesManager,
      final AccountsManager accountsManager,
      final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
      final Scheduler persistQueueScheduler,
      final Clock clock,
      final Duration persistDelay,
      final int maxConcurrency,
      final int scanCount,
      final RetryBackoffSpec retryBackoffSpec) {

    this.messagesCache = messagesCache;
    this.messagesManager = messagesManager;
    this.accountsManager = accountsManager;
    this.dynamicConfigurationManager = dynamicConfigurationManager;
    this.clock = clock;
    this.persistDelay = persistDelay;
    this.maxConcurrency = maxConcurrency;
    this.scanCount = scanCount;
    this.retryBackoffSpec = retryBackoffSpec;
    this.persistQueueScheduler = persistQueueScheduler;
  }

  @Override
  public synchronized void start() {
    running = true;

    workerThread = new Thread(() -> {
      while (running) {
        if (dynamicConfigurationManager.getConfiguration().getMessagePersisterConfiguration().isPersistenceEnabled()) {
          try {
            final int queuesPersisted = persistNextNode();
            QUEUE_COUNT_DISTRIBUTION_SUMMARY.record(queuesPersisted);

            if (queuesPersisted == 0) {
              Util.sleep(100);
            }
          } catch (final Throwable t) {
            logger.warn("Failed to persist queues", t);
            Util.sleep(EXCEPTION_PAUSE_MILLIS);
          }
        }

        try {
          // Persisters work through nodes quickly, and running in a hot loop taxes Redis nodes unnecessarily
          Thread.sleep(dynamicConfigurationManager.getConfiguration().getMessagePersisterConfiguration().getSleepBetweenNodes());
        } catch (final InterruptedException _) {
        }
      }
    }, "MessagePersister");

    workerThread.start();
  }

  @Override
  public synchronized void stop() {
    running = false;

    if (workerThread != null) {
      try {
        workerThread.join();
      } catch (final InterruptedException e) {
        logger.warn("Interrupted while waiting for worker thread to complete current operation");
      }

      workerThread = null;
    }
  }

  @VisibleForTesting
  int persistNextNode() {
    final RedisClusterNode node;
    {
      final Duration nodeClaimTtl =
          dynamicConfigurationManager.getConfiguration().getMessagePersisterConfiguration().getNodeClaimTtl();

      final Optional<RedisClusterNode> maybeNode = messagesCache.claimNextNodeToPersist(persisterId, nodeClaimTtl);

      if (maybeNode.isEmpty()) {
        return 0;
      }

      node = maybeNode.get();
    }

    try {
      return persistNode(node);
    } finally {
      messagesCache.releaseNodeClaim(node, persisterId);
    }
  }

  int persistNode(final RedisClusterNode node) {
    final LongTaskTimer.Sample sample = PERSIST_NODE_TIMER.start();

    try {
      final Tags tags = Tags.of("node", node.getUri().getHost());

      return messagesCache.getQueues(node, scanCount)
          .filter(_ -> dynamicConfigurationManager.getConfiguration().getMessagePersisterConfiguration()
              .isPersistenceEnabled())
          .flatMap(queueKey -> Mono.defer(() -> shouldPersistQueue(queueKey))
              .retryWhen(retryBackoffSpec)
              .onErrorResume(e -> {
                logger.warn("Failed to determine whether queue {} should be persisted", queueKey, e);
                return Mono.empty();
              })
              .mapNotNull(shouldPersist -> shouldPersist ? queueKey : null)
              .doOnTerminate(() -> Metrics.counter(INSPECTED_QUEUE_COUNTER_NAME, tags).increment()), maxConcurrency)
          .distinct()
          .flatMap(queueKey -> {
            final UUID accountIdentifier = MessagesCache.getAccountUuidFromQueueName(queueKey);
            final byte deviceId = MessagesCache.getDeviceIdFromQueueName(queueKey);

            return Mono.fromFuture(() -> accountsManager.getByAccountIdentifierAsync(accountIdentifier))
                .map(maybeAccount -> maybeAccount
                    .flatMap(account -> account.getDevice(deviceId))
                    .map(device -> Tuples.of(maybeAccount.get(), device)))
                .flatMap(Mono::justOrEmpty)
                .retryWhen(retryBackoffSpec)
                .onErrorResume(e -> {
                  logger.warn("Failed to fetch account/device for queue: {}", queueKey, e);
                  return Mono.empty();
                });
          }, maxConcurrency)
          .flatMap(accountAndDevice -> {
            final Account account = accountAndDevice.getT1();
            final Device device = accountAndDevice.getT2();

            return persistQueue(account, device, tags)
                .thenReturn(1)
                .retryWhen(retryBackoffSpec
                    // Don't retry with backoff for persistence exceptions
                    .filter(e -> !(e instanceof MessagePersistenceException)))
                .onErrorResume(e -> {
                  logger.warn("Failed to persist queue {}::{} ({}); will schedule for retry",
                      account.getIdentifier(IdentityType.ACI), device.getId(), node.getUri().getHost(), e);

                  return Mono.empty();
                });
          }, maxConcurrency)
          .onErrorResume(e -> {
            logger.warn("Unhandled error while persisting messages", e);
            return Mono.empty();
          })
          .count()
          .blockOptional()
          .map(Math::toIntExact)
          .orElse(0);
    } finally {
      sample.stop();
      PERSIST_NODE_COUNTER.increment();
    }
  }

  @VisibleForTesting
  Mono<Void> persistQueue(final Account account, final Device device, final Tags baseTags) {
    final UUID accountUuid = account.getUuid();
    final byte deviceId = device.getId();

    final Tag platformTag = Tag.of("platform", DevicePlatformUtil.getDevicePlatform(device)
        .map(platform -> platform.name().toLowerCase(Locale.ROOT))
        .orElse("unknown"));

    final Timer.Sample sample = Timer.start();
    final Tags tags = baseTags.and(platformTag);

    return Flux.usingWhen(
        messagesCache.lockQueueForPersistence(accountUuid, deviceId)
            .thenReturn(true),
        _ -> Flux.from(messagesCache.getMessagesToPersist(accountUuid, deviceId))
            .buffer(MESSAGE_BATCH_LIMIT)
            .flatMap(messages -> {
              final int urgentMessageCount = (int) messages.stream().filter(MessageProtos.Envelope::getUrgent).count();
              final int nonUrgentMessageCount = messages.size() - urgentMessageCount;

              Metrics.counter(PERSISTED_MESSAGE_COUNTER_NAME, tags.and("urgent", "true")).increment(urgentMessageCount);
              Metrics.counter(PERSISTED_MESSAGE_COUNTER_NAME, tags.and("urgent", "false")).increment(nonUrgentMessageCount);
              Metrics.counter(PERSISTED_BYTES_COUNTER_NAME, tags)
                  .increment(messages.stream().mapToInt(MessageProtos.Envelope::getSerializedSize).sum());

              return Mono.fromRunnable(() -> messagesManager.persistMessages(accountUuid, device, messages))
                  .subscribeOn(persistQueueScheduler)
                  .thenReturn(messages.size());
            }, 1)
            .reduce(0, Integer::sum)
            .onErrorResume(ItemCollectionSizeLimitExceededException.class, _ -> {
              final boolean isPrimary = deviceId == Device.PRIMARY_ID;
              Metrics.counter(OVERSIZED_QUEUE_COUNTER_NAME, "primary", String.valueOf(isPrimary)).increment();
              // may throw, in which case we'll retry later by the usual mechanism
              if (isPrimary) {
                logger.warn("Failed to persist queue {}::{} due to overfull queue; will trim oldest messages",
                    account.getUuid(), deviceId);

                return trimQueue(account, device)
                    .then(Mono.error(new MessagePersistenceException("Could not persist due to an overfull queue. Trimmed primary queue, a subsequent retry may succeed")));
              } else {
                logger.warn("Failed to persist queue {}::{} due to overfull queue; will unlink device", accountUuid, deviceId);

                return Mono.fromRunnable(() -> accountsManager.removeDevice(account, deviceId))
                    .subscribeOn(persistQueueScheduler)
                    .then(Mono.empty());
              }
            })
            .doOnSuccess(messagesPersisted -> {
              if (messagesPersisted != null) {
                DistributionSummary.builder(QUEUE_SIZE_DISTRIBUTION_SUMMARY_NAME)
                    .tags(Tags.of(platformTag))
                    .register(Metrics.globalRegistry)
                    .record(messagesPersisted);
              }
            })
            .doOnTerminate(() -> sample.stop(PERSIST_QUEUE_TIMER)),
        _ -> messagesCache.unlockQueueForPersistence(accountUuid, deviceId))
        .then();
  }

  private Mono<Void> trimQueue(final Account account, final Device device) {
    final UUID aci = account.getIdentifier(IdentityType.ACI);
    final byte deviceId = device.getId();

    final double extraRoomRatio = this.dynamicConfigurationManager.getConfiguration()
        .getMessagePersisterConfiguration()
        .getTrimOversizedQueueExtraRoomRatio();

    final AtomicLong oldestMessage = new AtomicLong(0L);
    final AtomicLong newestMessage = new AtomicLong(0L);
    final AtomicLong bytesDeleted = new AtomicLong(0L);

    final AtomicLong cachedMessageBytes = new AtomicLong(0L);
    final AtomicLong targetDeleteBytes = new AtomicLong(0L);

    return Mono.fromFuture(() -> messagesCache.estimatePersistedQueueSizeBytes(aci, deviceId))
        .flatMap(estimatedPersistedQueueSize -> {
          cachedMessageBytes.set(estimatedPersistedQueueSize);
          targetDeleteBytes.set(Math.round(estimatedPersistedQueueSize * extraRoomRatio));

          return Flux.from(messagesManager.getMessagesForDeviceReactive(aci, device, false))
              .concatMap(envelope -> {
                if (bytesDeleted.getAndAdd(envelope.getSerializedSize()) >= targetDeleteBytes.get()) {
                  return Mono.just(Optional.<MessageProtos.Envelope>empty());
                }
                oldestMessage.compareAndSet(0L, envelope.getServerTimestamp());
                newestMessage.set(envelope.getServerTimestamp());
                return Mono.just(Optional.of(envelope));
              })
              .takeWhile(Optional::isPresent)
              .flatMap(maybeEnvelope -> {
                // We know this must be present because we `takeWhile` values are present
                final MessageProtos.Envelope envelope = maybeEnvelope.orElseThrow(AssertionError::new);
                TRIMMED_MESSAGE_COUNTER.increment();
                TRIMMED_MESSAGE_BYTES_COUNTER.increment(envelope.getSerializedSize());
                return Mono
                    .fromCompletionStage(() -> messagesManager
                        .delete(aci, device, UUID.fromString(envelope.getServerGuid()), envelope.getServerTimestamp()))
                    .retryWhen(retryBackoffSpec)
                    .map(Optional::isPresent);
              })
              .reduce(Pair.of(0L, 0L), (acc, deleted) -> deleted
                  ? Pair.of(acc.getLeft() + 1, acc.getRight())
                  : Pair.of(acc.getLeft(), acc.getRight() + 1));
        })
        .doOnSuccess(outcomes -> {
          if (outcomes != null) {
            logger.warn(
                "Finished trimming {}:{}. Oldest message = {}, newest message = {}. Attempted to delete {} persisted bytes to make room for {} cached message bytes.  Delete outcomes: {} present, {} missing.",
                aci, deviceId,
                Instant.ofEpochMilli(oldestMessage.get()), Instant.ofEpochMilli(newestMessage.get()),
                targetDeleteBytes, cachedMessageBytes,
                outcomes.getLeft(), outcomes.getRight());
          }
        })
        .then();
  }

  @VisibleForTesting
  Mono<Boolean> shouldPersistQueue(final String queueKey) {
    return messagesCache.getEarliestUndeliveredTimestamp(
            MessagesCache.getAccountUuidFromQueueName(queueKey),
            MessagesCache.getDeviceIdFromQueueName(queueKey))
        .map(oldestMessageTimestamp ->
            Duration.between(Instant.ofEpochMilli(oldestMessageTimestamp), clock.instant()).compareTo(persistDelay) > 0)
        .switchIfEmpty(Mono.just(false));
  }

  @VisibleForTesting
  String getPersisterId() {
    return persisterId;
  }
}
