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

import java.util.concurrent.Callable;
import org.apache.beam.it.common.ResourceManager;
import org.apache.beam.it.common.TestProperties;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.containers.GenericContainer;

/**
 * Abstract class for managing TestContainers resources in integration tests.
 *
 * <p>Resource managers that extend this class will use TestContainers as a backend to spin up
 * resources.
 *
 * <p>Optionally, a static resource can be specified by calling the useStaticContainer() method in
 * the {@link Builder} class. A static resource is a pre-configured database or other resource that
 * is ready to be connected to by the resource manager. This could be a pre-existing TestContainer
 * that has not been closed, a local database instance, a remote VM, or any other source that can be
 * connected to. If a static container is used, the host and port must also be configured using the
 * Builder's setHost() and setPort() methods, respectively.
 */
public abstract class TestContainerResourceManager<T extends GenericContainer<?>>
    implements ResourceManager {
  private static final Logger LOG = LoggerFactory.getLogger(TestContainerResourceManager.class);

  protected static final String HOST_IP = TestProperties.hostIp();
  private final T container;
  private final boolean usingStaticContainer;
  private final String host;
  protected int port;

  protected <B extends Builder<?>> TestContainerResourceManager(T container, B builder) {
    this(container, builder, null);
  }

  protected <B extends Builder<?>> TestContainerResourceManager(
      T container, B builder, @Nullable Callable<Void> setup) {
    this.container = container;
    this.usingStaticContainer = builder.useStaticContainer;
    this.host = builder.host == null ? HOST_IP : builder.host;
    this.port = builder.port;

    if (setup != null) {
      try {
        setup.call();
      } catch (Exception e) {
        throw new TestContainerResourceManagerException("Error running setup function.", e);
      }
    }

    if (!usingStaticContainer) {
      // TODO(pranavbhandari): Change this to use log.getUtf8StringWithoutLineEnding() when
      // testcontainers dependency is updated.
      container
          .withLogConsumer(
              log -> LOG.info("{}: {}", container.getDockerImageName(), log.getUtf8String()))
          .start();
    } else if (builder.host == null || builder.port < 0) {
      throw new TestContainerResourceManagerException(
          "This manager was configured to use a static resource, but the host and port were not properly set.");
    }
  }

  protected String getDockerImageName() {
    return container.getDockerImageName();
  }

  /**
   * Returns the host that the resource is running on. If a host was specified in the builder, then
   * this method will return that value. Otherwise, the host will default to localhost.
   *
   * @return the host that the resource is running on
   */
  public String getHost() {
    return host;
  }

  /**
   * Returns the port that the resource is actually mapped to on the host. When using a non-static
   * resource, the TestContainers framework will map the resource port to a random port on the host.
   * This method will return the mapped port that the resource traffic is routed through.
   *
   * <p>Note: When using a static container, the port that was specified in the builder will be
   * returned regardless of what is passed into this method as a parameter.
   *
   * @param mappedPort the port the resource was configured to listen on
   * @return the mapped port on the host
   */
  protected int getPort(int mappedPort) {
    if (port < 0) {
      return container.getMappedPort(mappedPort);
    }
    return port;
  }

  /**
   * Deletes all created resources (VM's, etc.) and stops the container, making the manager object
   * unusable.
   *
   * @throws TestContainerResourceManagerException if there is an error deleting the TestContainers
   *     resources.
   */
  @Override
  public void cleanupAll() throws TestContainerResourceManagerException {
    if (usingStaticContainer) {
      LOG.info("This manager was configured to use a static resource that will not be cleaned up.");
      return;
    }

    LOG.info("Attempting to cleanup TestContainers manager.");
    try {
      container.close();
      LOG.info("TestContainers manager successfully cleaned up.");
    } catch (Exception e) {
      throw new TestContainerResourceManagerException("Failed to close TestContainer resources", e);
    }
  }

  /** Builder for {@link TestContainerResourceManager}. */
  public abstract static class Builder<T extends TestContainerResourceManager<?>> {

    public String testId;
    public String containerImageName;
    public String containerImageTag;
    public @Nullable String host;
    public int port;
    public boolean useStaticContainer;

    public Builder(String testId, String containerImageName, String containerImageTag) {
      this.testId = testId;
      this.containerImageName = containerImageName;
      this.containerImageTag = containerImageTag;
      this.host = null;
      this.port = -1;
    }

    /**
     * Sets the name of the test container image. The tag is typically the version of the image.
     *
     * @param containerName The name of the container image.
     * @return this builder object with the image name set.
     */
    public Builder<T> setContainerImageName(String containerName) {
      this.containerImageName = containerName;
      return this;
    }

    /**
     * Sets the tag for the test container. The tag is typically the version of the image.
     *
     * @param containerTag The tag to use for the container.
     * @return this builder object with the tag set.
     */
    public Builder<T> setContainerImageTag(String containerTag) {
      this.containerImageTag = containerTag;
      return this;
    }

    /**
     * Sets the host of the resource that the resource manager will connect to. This will default to
     * localhost if not set.
     *
     * @param containerHost the resource host address.
     * @return this builder object with the host set.
     */
    public Builder<T> setHost(String containerHost) {
      this.host = containerHost;
      return this;
    }

    /**
     * Sets the port that the resource is hosted on.
     *
     * @param port the port the resource is hosted on.
     * @return this builder object with the port set.
     */
    public Builder<T> setPort(int port) {
      this.port = port;
      return this;
    }

    /**
     * Configures the resource manager to use a static resource instead of creating a new
     * TestContainer instance of the resource. This can be another TestContainer or any other VM or
     * server hosting the resource.
     *
     * <p>Note: When this option is enabled, the setPort() and setHost() methods must also be called
     * to configure the static resource address.
     *
     * @return this builder object with the useStaticContainer option enabled.
     */
    public Builder<T> useStaticContainer() {
      this.useStaticContainer = true;
      return this;
    }

    /**
     * Builds and returns a Resource Manager that extends TestContainerResourceManager.
     *
     * @return an instance of TestContainerResourceManager
     */
    public abstract T build();
  }
}
