/*
 * 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.it.mongodb;

import static org.apache.beam.it.mongodb.MongoDBResourceManagerUtils.checkValidCollectionName;
import static org.apache.beam.it.mongodb.MongoDBResourceManagerUtils.generateDatabaseName;

import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoClient;
import com.mongodb.client.MongoClients;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import java.util.List;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;
import org.apache.beam.it.common.ResourceManager;
import org.apache.beam.it.testcontainers.TestContainerResourceManager;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.bson.Document;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.containers.MongoDBContainer;
import org.testcontainers.utility.DockerImageName;

/**
 * Client for managing MongoDB resources.
 *
 * <p>The class supports one database and multiple collections per database object. A database is
 * created when the first collection is created if one has not been created already.
 *
 * <p>The database name is formed using testId. The database name will be "{testId}-{ISO8601 time,
 * microsecond precision}", with additional formatting.
 *
 * <p>The class is thread-safe.
 */
public class MongoDBResourceManager extends TestContainerResourceManager<MongoDBContainer>
    implements ResourceManager {

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

  private static final String DEFAULT_MONGODB_CONTAINER_NAME = "mongo";

  // A list of available MongoDB Docker image tags can be found at
  // https://hub.docker.com/_/mongo/tags
  private static final String DEFAULT_MONGODB_CONTAINER_TAG = "4.0.18";

  // 27017 is the default port that MongoDB is configured to listen on
  private static final int MONGODB_INTERNAL_PORT = 27017;

  private final MongoClient mongoClient;
  private final String databaseName;
  private final String connectionString;
  private final boolean usingStaticDatabase;

  private MongoDBResourceManager(Builder builder) {
    this(
        /* mongoClient= */ null,
        new MongoDBContainer(
            DockerImageName.parse(builder.containerImageName).withTag(builder.containerImageTag)),
        builder);
  }

  @VisibleForTesting
  @SuppressWarnings("nullness")
  MongoDBResourceManager(
      @Nullable MongoClient mongoClient, MongoDBContainer container, Builder builder) {
    super(container, builder);

    this.usingStaticDatabase = builder.databaseName != null;
    this.databaseName =
        usingStaticDatabase ? builder.databaseName : generateDatabaseName(builder.testId);
    this.connectionString =
        String.format("mongodb://%s:%d", this.getHost(), this.getPort(MONGODB_INTERNAL_PORT));
    this.mongoClient = mongoClient == null ? MongoClients.create(connectionString) : mongoClient;
  }

  public static Builder builder(String testId) {
    return new Builder(testId);
  }

  /** Returns the URI connection string to the MongoDB Database. */
  public synchronized String getUri() {
    return connectionString;
  }

  /**
   * Returns the name of the Database that this MongoDB manager will operate in.
   *
   * @return the name of the MongoDB Database.
   */
  public synchronized String getDatabaseName() {
    return databaseName;
  }

  private synchronized MongoDatabase getDatabase() {
    try {
      return mongoClient.getDatabase(databaseName);
    } catch (Exception e) {
      throw new MongoDBResourceManagerException(
          "Error retrieving database " + databaseName + " from MongoDB.", e);
    }
  }

  private synchronized boolean collectionExists(String collectionName) {
    // Check collection name
    checkValidCollectionName(databaseName, collectionName);

    Iterable<String> collectionNames = getDatabase().listCollectionNames();
    for (String name : collectionNames) {
      // The Collection already exists in the database, return false.
      if (collectionName.equals(name)) {
        return true;
      }
    }

    return false;
  }

  /**
   * Creates a collection in MongoDB for storing Documents.
   *
   * <p>Note: Implementations may do database creation here, if one does not already exist.
   *
   * @param collectionName Collection name to associate with the given MongoDB instance.
   * @return A boolean indicating whether the resource was created.
   * @throws MongoDBResourceManagerException if there is an error creating the collection in
   *     MongoDB.
   */
  public synchronized boolean createCollection(String collectionName)
      throws MongoDBResourceManagerException {
    LOG.info("Creating collection using collectionName '{}'.", collectionName);

    try {
      // Check to see if the Collection exists
      if (collectionExists(collectionName)) {
        return false;
      }
      // The Collection does not exist in the database, create it and return true.
      getDatabase().createCollection(collectionName);
    } catch (Exception e) {
      throw new MongoDBResourceManagerException("Error creating collection.", e);
    }

    LOG.info("Successfully created collection {}.{}", databaseName, collectionName);

    return true;
  }

  /**
   * Helper method to retrieve a MongoCollection with the given name from the database and return
   * it.
   *
   * @param collectionName The name of the MongoCollection.
   * @param createCollection A boolean that specifies to create the Collection if it does not exist.
   * @return A MongoCollection with the given name.
   */
  private MongoCollection<Document> getMongoDBCollection(
      String collectionName, boolean createCollection) {
    if (!collectionExists(collectionName) && !createCollection) {
      throw new MongoDBResourceManagerException(
          "Collection " + collectionName + " does not exists in database " + databaseName);
    }

    return getDatabase().getCollection(collectionName);
  }

  /**
   * Inserts the given Document into a collection.
   *
   * <p>A database will be created here, if one does not already exist.
   *
   * @param collectionName The name of the collection to insert the document into.
   * @param document The document to insert into the collection.
   * @return A boolean indicating whether the Document was inserted successfully.
   */
  public synchronized boolean insertDocument(String collectionName, Document document) {
    return insertDocuments(collectionName, ImmutableList.of(document));
  }

  /**
   * Inserts the given Documents into a collection.
   *
   * <p>Note: Implementations may do collection creation here, if one does not already exist.
   *
   * @param collectionName The name of the collection to insert the documents into.
   * @param documents A list of documents to insert into the collection.
   * @return A boolean indicating whether the Documents were inserted successfully.
   * @throws MongoDBResourceManagerException if there is an error inserting the documents.
   */
  public synchronized boolean insertDocuments(String collectionName, List<Document> documents)
      throws MongoDBResourceManagerException {
    LOG.info(
        "Attempting to write {} documents to {}.{}.",
        documents.size(),
        databaseName,
        collectionName);

    try {
      getMongoDBCollection(collectionName, /* createCollection= */ true).insertMany(documents);
    } catch (Exception e) {
      throw new MongoDBResourceManagerException("Error inserting documents.", e);
    }

    LOG.info(
        "Successfully wrote {} documents to {}.{}", documents.size(), databaseName, collectionName);

    return true;
  }

  /**
   * Reads all the Documents in a collection.
   *
   * @param collectionName The name of the collection to read from.
   * @return A List of all the Documents in the collection.
   * @throws MongoDBResourceManagerException if there is an error reading the collection.
   */
  public synchronized List<Document> readCollection(String collectionName)
      throws MongoDBResourceManagerException {
    LOG.info("Reading all documents from {}.{}", databaseName, collectionName);

    FindIterable<Document> fetchRecords;
    try {
      fetchRecords = getMongoDBCollection(collectionName, /* createCollection= */ false).find();
    } catch (Exception e) {
      throw new MongoDBResourceManagerException("Error reading collection.", e);
    }
    List<Document> documents =
        StreamSupport.stream(fetchRecords.spliterator(), false).collect(Collectors.toList());

    LOG.info("Successfully loaded documents from {}.{}", databaseName, collectionName);

    return documents;
  }

  /**
   * Counts the number of Documents in a collection.
   *
   * @param collectionName The name of the collection to read from.
   * @return The number of Documents in the collection.
   * @throws MongoDBResourceManagerException if there is an error reading the collection.
   */
  public synchronized long countCollection(String collectionName)
      throws MongoDBResourceManagerException {
    LOG.info("Counting all documents from {}.{}", databaseName, collectionName);

    long numDocuments =
        getMongoDBCollection(collectionName, /* createCollection= */ false).countDocuments();

    LOG.info(
        "Successfully counted {} documents from {}.{}", numDocuments, databaseName, collectionName);

    return numDocuments;
  }

  @Override
  public synchronized void cleanupAll() {
    LOG.info("Attempting to cleanup MongoDB manager.");

    boolean producedError = false;

    // First, delete the database if it was not given as a static argument
    try {
      if (!usingStaticDatabase) {
        mongoClient.getDatabase(databaseName).drop();
      }
    } catch (Exception e) {
      LOG.error("Failed to delete MongoDB database {}.", databaseName, e);
      producedError = true;
    }

    // Next, try to close the MongoDB client connection
    try {
      mongoClient.close();
    } catch (Exception e) {
      LOG.error("Failed to delete MongoDB client.", e);
      producedError = true;
    }

    // Throw Exception at the end if there were any errors
    if (producedError) {
      throw new MongoDBResourceManagerException(
          "Failed to delete resources. Check above for errors.");
    }

    super.cleanupAll();

    LOG.info("MongoDB manager successfully cleaned up.");
  }

  /** Builder for {@link MongoDBResourceManager}. */
  public static final class Builder
      extends TestContainerResourceManager.Builder<MongoDBResourceManager> {

    private @Nullable String databaseName;

    private Builder(String testId) {
      super(testId, DEFAULT_MONGODB_CONTAINER_NAME, DEFAULT_MONGODB_CONTAINER_TAG);
      this.databaseName = null;
    }

    /**
     * Sets the database name to that of a static database instance. Use this method only when
     * attempting to operate on a pre-existing Mongo database.
     *
     * <p>Note: if a database name is set, and a static MongoDB server is being used
     * (useStaticContainer() is also called on the builder), then a database will be created on the
     * static server if it does not exist, and it will not be removed when cleanupAll() is called on
     * the MongoDBResourceManager.
     *
     * @param databaseName The database name.
     * @return this builder object with the database name set.
     */
    public Builder setDatabaseName(String databaseName) {
      this.databaseName = databaseName;
      return this;
    }

    @Override
    public MongoDBResourceManager build() {
      return new MongoDBResourceManager(this);
    }
  }
}
