/*
 * 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.beam.sdk.nexmark;

import com.google.api.client.googleapis.json.GoogleJsonResponseException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.SubscriptionPath;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubClient.TopicPath;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubJsonClient;
import org.apache.beam.sdk.io.gcp.pubsub.PubsubOptions;
import org.apache.beam.sdk.util.BackOff;
import org.apache.beam.sdk.util.BackOffUtils;
import org.apache.beam.sdk.util.FluentBackoff;
import org.apache.beam.sdk.util.Sleeper;
import org.joda.time.Duration;

/** Helper for working with pubsub and gcs. */
public class PubsubHelper {
  /** Underlying pub/sub client. */
  private final PubsubClient pubsubClient;

  /** Project id. */
  private final String project;

  /** Topics we should delete on close. */
  private final List<TopicPath> createdTopics;

  /** Subscriptions we should delete on close. */
  private final List<SubscriptionPath> createdSubscriptions;

  /** How to sleep in retry loops. */
  private final Sleeper sleeper;

  /** How to backoff in retry loops. */
  private final BackOff backOff;

  private PubsubHelper(PubsubClient pubsubClient, String project) {
    this.pubsubClient = pubsubClient;
    this.project = project;
    createdTopics = new ArrayList<>();
    createdSubscriptions = new ArrayList<>();
    sleeper = Sleeper.DEFAULT;
    backOff =
        FluentBackoff.DEFAULT
            .withInitialBackoff(Duration.standardSeconds(1))
            .withMaxRetries(3)
            .backoff();
  }

  /** Create a helper. */
  public static PubsubHelper create(PubsubOptions options) throws IOException {
    return new PubsubHelper(
        PubsubJsonClient.FACTORY.newClient(null, null, options), options.getProject());
  }

  /**
   * Create a topic from short name. Delete it if it already exists. Ensure the topic will be
   * deleted on cleanup. Return full topic name.
   */
  public TopicPath createTopic(String shortTopic) throws IOException {
    TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
    while (true) {
      try {
        NexmarkUtils.console("create topic %s", topic);
        pubsubClient.createTopic(topic);
        createdTopics.add(topic);
        return topic;
      } catch (GoogleJsonResponseException ex) {
        NexmarkUtils.console("attempting to cleanup topic %s", topic);
        pubsubClient.deleteTopic(topic);
        try {
          if (!BackOffUtils.next(sleeper, backOff)) {
            NexmarkUtils.console("too many retries for creating topic %s", topic);
            throw ex;
          }
        } catch (InterruptedException in) {
          throw new IOException(in);
        }
      }
    }
  }

  /**
   * Create a topic from short name if it does not already exist. The topic will not be deleted on
   * cleanup. Return full topic name.
   */
  public TopicPath createOrReuseTopic(String shortTopic) throws IOException {
    TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
    while (true) {
      try {
        NexmarkUtils.console("create topic %s", topic);
        pubsubClient.createTopic(topic);
        return topic;
      } catch (GoogleJsonResponseException ex) {
        if (topicExists(shortTopic)) {
          NexmarkUtils.console("topic %s already exists", topic);
          return topic;
        }
        try {
          if (!BackOffUtils.next(sleeper, backOff)) {
            NexmarkUtils.console("too many retries for creating/reusing topic %s", topic);
            throw ex;
          }
        } catch (InterruptedException in) {
          throw new IOException(in);
        }
      }
    }
  }

  /**
   * Check a topic corresponding to short name exists, and throw exception if not. The topic will
   * not be deleted on cleanup. Return full topic name.
   */
  public TopicPath reuseTopic(String shortTopic) throws IOException {
    TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
    if (topicExists(shortTopic)) {
      NexmarkUtils.console("reusing existing topic %s", topic);
      return topic;
    }
    throw new RuntimeException("topic '" + topic + "' does not already exist");
  }

  /** Does topic corresponding to short name exist? */
  public boolean topicExists(String shortTopic) throws IOException {
    TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
    return pubsubClient
        .listTopics(PubsubClient.projectPathFromId(project))
        .stream()
        .anyMatch(topic::equals);
  }

  /**
   * Create subscription from short name. Ensure the subscription will be deleted on cleanup. Return
   * full subscription name.
   */
  public SubscriptionPath createSubscription(String shortTopic, String shortSubscription)
      throws IOException {
    TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
    SubscriptionPath subscription =
        PubsubClient.subscriptionPathFromName(project, shortSubscription);
    while (true) {
      try {
        NexmarkUtils.console("create subscription %s", subscription);
        pubsubClient.createSubscription(topic, subscription, 60);
        createdSubscriptions.add(subscription);
        return subscription;
      } catch (GoogleJsonResponseException ex) {
        NexmarkUtils.console("attempting to cleanup subscription %s", subscription);
        pubsubClient.deleteSubscription(subscription);
        try {
          if (!BackOffUtils.next(sleeper, backOff)) {
            NexmarkUtils.console("too many retries for creating subscription %s", subscription);
            throw ex;
          }
        } catch (InterruptedException in) {
          throw new IOException(in);
        }
      }
    }
  }

  /**
   * Check a subscription corresponding to short name exists, and throw exception if not. The
   * subscription will not be deleted on cleanup. Return full topic name.
   */
  public SubscriptionPath reuseSubscription(String shortTopic, String shortSubscription)
      throws IOException {
    SubscriptionPath subscription =
        PubsubClient.subscriptionPathFromName(project, shortSubscription);
    if (subscriptionExists(shortTopic, shortSubscription)) {
      NexmarkUtils.console("reusing existing subscription %s", subscription);
      return subscription;
    }
    throw new RuntimeException("subscription'" + subscription + "' does not already exist");
  }

  /** Does subscription corresponding to short name exist? */
  public boolean subscriptionExists(String shortTopic, String shortSubscription)
      throws IOException {
    TopicPath topic = PubsubClient.topicPathFromName(project, shortTopic);
    SubscriptionPath subscription =
        PubsubClient.subscriptionPathFromName(project, shortSubscription);
    return pubsubClient
        .listSubscriptions(PubsubClient.projectPathFromId(project), topic)
        .stream()
        .anyMatch(subscription::equals);
  }

  /** Delete all the subscriptions and topics we created. */
  public void cleanup() {
    for (SubscriptionPath subscription : createdSubscriptions) {
      try {
        NexmarkUtils.console("delete subscription %s", subscription);
        pubsubClient.deleteSubscription(subscription);
      } catch (IOException ex) {
        NexmarkUtils.console("could not delete subscription %s", subscription);
      }
    }
    for (TopicPath topic : createdTopics) {
      try {
        NexmarkUtils.console("delete topic %s", topic);
        pubsubClient.deleteTopic(topic);
      } catch (IOException ex) {
        NexmarkUtils.console("could not delete topic %s", topic);
      }
    }
  }
}
