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

import static org.apache.beam.it.gcp.LoadTestBase.createConfig;
import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatPipeline;
import static org.apache.beam.it.truthmatchers.PipelineAsserts.assertThatResult;

import com.google.auth.Credentials;
import java.io.IOException;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;
import org.apache.beam.it.common.PipelineLauncher;
import org.apache.beam.it.common.PipelineLauncher.LaunchConfig;
import org.apache.beam.it.common.PipelineLauncher.LaunchInfo;
import org.apache.beam.it.common.PipelineOperator;
import org.apache.beam.it.common.PipelineOperator.Result;
import org.apache.beam.it.common.TestProperties;
import org.apache.beam.it.gcp.dataflow.FlexTemplateClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/** Helper class for starting a Streaming Data Generator Dataflow template job. */
public class DataGenerator {
  private static final Logger LOG = LoggerFactory.getLogger(DataGenerator.class);
  private static final String SPEC_PATH =
      "gs://dataflow-templates/latest/flex/Streaming_Data_Generator";
  private static final String PROJECT = TestProperties.project();
  private static final String REGION = TestProperties.region();
  private static final Credentials CREDENTIALS = TestProperties.googleCredentials();
  private static final String MESSAGES_GENERATED_METRIC_NAME =
      "Generate Fake Messages-out0-ElementCount";
  private static final String MESSAGES_LIMIT = "messagesLimit";
  private final LaunchConfig dataGeneratorOptions;
  private final PipelineLauncher pipelineLauncher;
  private final PipelineOperator pipelineOperator;

  private DataGenerator(Builder builder) {
    pipelineLauncher = FlexTemplateClient.builder(CREDENTIALS).build();
    pipelineOperator = new PipelineOperator(pipelineLauncher);
    this.dataGeneratorOptions =
        LaunchConfig.builder(builder.getJobName(), SPEC_PATH)
            .setParameters(builder.getParameters())
            .addParameter("experiments", "disable_runner_v2")
            .build();
  }

  public static Builder builderWithSchemaLocation(String testName, String schemaLocation) {
    return new Builder(testName + "-data-generator")
        .setSchemaLocation(schemaLocation)
        .setAutoscalingAlgorithm(AutoscalingAlgorithmType.THROUGHPUT_BASED);
  }

  public static Builder builderWithSchemaTemplate(String testName, String schemaTemplate) {
    return new Builder(testName + "-data-generator")
        .setSchemaTemplate(schemaTemplate)
        .setAutoscalingAlgorithm(AutoscalingAlgorithmType.THROUGHPUT_BASED);
  }

  /**
   * Executes the data generator using the config provided. If messageLimit is provided, we wait
   * until the data generator finishes, or we reach timeout. If a messageLimit is not provided we
   * wait until timeout and cancel the data generator.
   *
   * <p>Note: This is a blocking call. For backlog tests, start the Data generator before executing
   * the template under test. For testing a pipeline against live incoming data, execute the data
   * generator after starting the pipeline under test.
   *
   * @param timeout time to wait before cancelling the data generator.
   * @return approximate number of messages generated
   * @throws IOException if any errors are encountered.
   */
  public Integer execute(Duration timeout) throws IOException {
    LaunchInfo dataGeneratorLaunchInfo =
        pipelineLauncher.launch(PROJECT, REGION, dataGeneratorOptions);
    assertThatPipeline(dataGeneratorLaunchInfo).isRunning();
    PipelineOperator.Config config = createConfig(dataGeneratorLaunchInfo, timeout);
    // check if the job will be BATCH or STREAMING
    if (dataGeneratorOptions.parameters().containsKey(MESSAGES_LIMIT)) {
      // Batch job, wait till data generator job finishes
      Result dataGeneratorResult = pipelineOperator.waitUntilDone(config);
      assertThatResult(dataGeneratorResult).isLaunchFinished();
    } else {
      // Streaming job, wait till timeout and drain job
      Result dataGeneratorResult = pipelineOperator.waitUntilDoneAndFinish(config);
      assertThatResult(dataGeneratorResult).hasTimedOut();
    }
    @SuppressWarnings("nullness")
    int generatedMessages =
        pipelineLauncher
            .getMetric(
                PROJECT, REGION, dataGeneratorLaunchInfo.jobId(), MESSAGES_GENERATED_METRIC_NAME)
            .intValue();
    LOG.info("Data generator finished. Generated {} messages.", generatedMessages);
    return generatedMessages;
  }

  /** Builder for the {@link DataGenerator}. */
  public static final class Builder {
    private final String jobName;
    private final Map<String, String> parameters;

    private Builder(String jobName) {
      this.jobName = jobName;
      this.parameters = new HashMap<>();
    }

    public String getJobName() {
      return jobName;
    }

    public Map<String, String> getParameters() {
      return parameters;
    }

    public Builder setSchemaTemplate(String value) {
      parameters.put("schemaTemplate", value);
      return this;
    }

    public Builder setSchemaLocation(String value) {
      parameters.put("schemaLocation", value);
      return this;
    }

    public Builder setMessagesLimit(String value) {
      parameters.put(MESSAGES_LIMIT, value);
      return this;
    }

    public Builder setQPS(String value) {
      parameters.put("qps", value);
      return this;
    }

    public Builder setSinkType(String value) {
      parameters.put("sinkType", value);
      return this;
    }

    public Builder setWorkerMachineType(String value) {
      parameters.put("workerMachineType", value);
      return this;
    }

    public Builder setNumWorkers(String value) {
      parameters.put("numWorkers", value);
      return this;
    }

    public Builder setMaxNumWorkers(String value) {
      parameters.put("maxNumWorkers", value);
      return this;
    }

    public Builder setAutoscalingAlgorithm(AutoscalingAlgorithmType value) {
      parameters.put("autoscalingAlgorithm", value.toString());
      return this;
    }

    public Builder setOutputDirectory(String value) {
      parameters.put("outputDirectory", value);
      return this;
    }

    public Builder setOutputType(String value) {
      parameters.put("outputType", value);
      return this;
    }

    public Builder setNumShards(String value) {
      parameters.put("numShards", value);
      return this;
    }

    public Builder setAvroSchemaLocation(String value) {
      parameters.put("avroSchemaLocation", value);
      return this;
    }

    public Builder setTopic(String value) {
      parameters.put("topic", value);
      return this;
    }

    public Builder setProjectId(String value) {
      parameters.put("projectId", value);
      return this;
    }

    public Builder setSpannerInstanceName(String value) {
      parameters.put("spannerInstanceName", value);
      return this;
    }

    public Builder setSpannerDatabaseName(String value) {
      parameters.put("spannerDatabaseName", value);
      return this;
    }

    public Builder setSpannerTableName(String value) {
      parameters.put("spannerTableName", value);
      return this;
    }

    public Builder setDriverClassName(String value) {
      parameters.put("driverClassName", value);
      return this;
    }

    public Builder setConnectionUrl(String value) {
      parameters.put("connectionUrl", value);
      return this;
    }

    public Builder setUsername(String value) {
      parameters.put("username", value);
      return this;
    }

    public Builder setPassword(String value) {
      parameters.put("password", value);
      return this;
    }

    public Builder setConnectionProperties(String value) {
      parameters.put("connectionProperties", value);
      return this;
    }

    public Builder setStatement(String value) {
      parameters.put("statement", value);
      return this;
    }

    public DataGenerator build() {
      return new DataGenerator(this);
    }
  }

  /** Enum representing Autoscaling algorithm types. */
  public enum AutoscalingAlgorithmType {
    NONE("NONE"),
    THROUGHPUT_BASED("THROUGHPUT_BASED");

    private final String text;

    AutoscalingAlgorithmType(String text) {
      this.text = text;
    }

    @Override
    public String toString() {
      return this.text;
    }
  }
}
