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

import static org.apache.beam.it.gcp.bigtable.BigtableResourceManagerUtils.generateTableId;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;

import com.google.auto.value.AutoValue;
import com.google.bigtable.v2.Mutation;
import com.google.protobuf.ByteString;
import java.io.IOException;
import java.io.Serializable;
import java.text.ParseException;
import java.time.Duration;
import java.time.Instant;
import java.util.Map;
import java.util.Random;
import org.apache.beam.it.common.PipelineLauncher;
import org.apache.beam.it.common.PipelineOperator;
import org.apache.beam.it.common.TestProperties;
import org.apache.beam.it.common.utils.ResourceManagerUtils;
import org.apache.beam.it.gcp.IOLoadTestBase;
import org.apache.beam.sdk.io.GenerateSequence;
import org.apache.beam.sdk.io.gcp.bigtable.BigtableIO;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.testing.TestPipelineOptions;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;

/**
 * BigTableIO performance tests.
 *
 * <p>Example trigger command for all tests: "mvn test -pl it/google-cloud-platform -am
 * -Dtest=BigTableIOLT \ -Dproject=[gcpProject] -DartifactBucket=[temp bucket]
 * -DfailIfNoTests=false".
 *
 * <p>Example trigger command for specific test: "mvn test -pl it/google-cloud-platform -am \
 * -Dtest="BigTableIOLT#testBigtableWriteAndRead" -Dconfiguration=local -Dproject=[gcpProject] \
 * -DartifactBucket=[temp bucket] -DfailIfNoTests=false".
 */
public class BigTableIOLT extends IOLoadTestBase {

  private static final String COLUMN_FAMILY_NAME = "cf";
  private static final long TABLE_MAX_AGE_MINUTES = 100L;

  private BigtableResourceManager resourceManager;
  private static final String READ_ELEMENT_METRIC_NAME = "read_count";
  private Configuration configuration;
  private String tableId;

  @Rule public TestPipeline writePipeline = TestPipeline.create();
  @Rule public TestPipeline readPipeline = TestPipeline.create();

  @Before
  public void setup() throws IOException {
    resourceManager =
        BigtableResourceManager.builder(testName, project, CREDENTIALS_PROVIDER).build();

    String testConfig =
        TestProperties.getProperty("configuration", "small", TestProperties.Type.PROPERTY);
    configuration = TEST_CONFIGS.get(testConfig);
    if (configuration == null) {
      throw new IllegalArgumentException(
          String.format(
              "Unknown test configuration: [%s]. Known configs: %s",
              testConfig, TEST_CONFIGS.keySet()));
    }
    // tempLocation needs to be set for bigtable IO writes
    if (!Strings.isNullOrEmpty(tempBucketName)) {
      String tempLocation = String.format("gs://%s/temp/", tempBucketName);
      writePipeline.getOptions().as(TestPipelineOptions.class).setTempRoot(tempLocation);
      writePipeline.getOptions().setTempLocation(tempLocation);
      readPipeline.getOptions().as(TestPipelineOptions.class).setTempRoot(tempLocation);
      readPipeline.getOptions().setTempLocation(tempLocation);
    }
  }

  @After
  public void teardown() {
    ResourceManagerUtils.cleanResources(resourceManager);
  }

  private static final Map<String, Configuration> TEST_CONFIGS =
      ImmutableMap.of(
          "local", Configuration.of(1_000L, 5, "DirectRunner", 1000), // 1MB
          "small", Configuration.of(1_000_000L, 20, "DataflowRunner", 1000), // 1 GB
          "medium", Configuration.of(10_000_000L, 40, "DataflowRunner", 1000), // 10 GB
          "large", Configuration.of(100_000_000L, 100, "DataflowRunner", 1000) // 100 GB
          );

  /** Run integration test with configurations specified by TestProperties. */
  @Test
  public void testBigtableWriteAndRead() throws IOException {

    tableId = generateTableId(testName);
    resourceManager.createTable(
        tableId,
        ImmutableList.of(COLUMN_FAMILY_NAME),
        org.threeten.bp.Duration.ofMinutes(TABLE_MAX_AGE_MINUTES));

    PipelineLauncher.LaunchInfo writeInfo = testWrite();
    PipelineOperator.Result writeResult =
        pipelineOperator.waitUntilDone(
            createConfig(writeInfo, Duration.ofMinutes(configuration.getPipelineTimeout())));
    assertNotEquals(PipelineOperator.Result.LAUNCH_FAILED, writeResult);

    PipelineLauncher.LaunchInfo readInfo = testRead();
    PipelineOperator.Result result =
        pipelineOperator.waitUntilDone(
            createConfig(readInfo, Duration.ofMinutes(configuration.getPipelineTimeout())));
    assertNotEquals(PipelineOperator.Result.LAUNCH_FAILED, result);
    assertEquals(
        PipelineLauncher.JobState.DONE,
        pipelineLauncher.getJobStatus(project, region, readInfo.jobId()));
    double numRecords =
        pipelineLauncher.getMetric(
            project,
            region,
            readInfo.jobId(),
            getBeamMetricsName(PipelineMetricsType.COUNTER, READ_ELEMENT_METRIC_NAME));
    assertEquals(configuration.getNumRows(), numRecords, 0.5);

    // export metrics
    MetricsConfiguration metricsConfig =
        MetricsConfiguration.builder()
            .setInputPCollection("Map records.out0")
            .setInputPCollectionV2("Map records/ParMultiDo(MapToBigTableFormat).out0")
            .setOutputPCollection("Counting element.out0")
            .setOutputPCollectionV2("Counting element/ParMultiDo(Counting).out0")
            .build();
    try {
      exportMetricsToBigQuery(writeInfo, getMetrics(writeInfo, metricsConfig));
      exportMetricsToBigQuery(readInfo, getMetrics(readInfo, metricsConfig));
    } catch (ParseException | InterruptedException e) {
      throw new RuntimeException(e);
    }
  }

  private PipelineLauncher.LaunchInfo testWrite() throws IOException {

    BigtableIO.Write writeIO =
        BigtableIO.write()
            .withProjectId(project)
            .withInstanceId(resourceManager.getInstanceId())
            .withTableId(tableId);

    writePipeline
        .apply(GenerateSequence.from(0).to(configuration.getNumRows()))
        .apply("Map records", ParDo.of(new MapToBigTableFormat(configuration.getValueSizeBytes())))
        .apply("Write to BigTable", writeIO);

    PipelineLauncher.LaunchConfig options =
        PipelineLauncher.LaunchConfig.builder("write-bigtable")
            .setSdk(PipelineLauncher.Sdk.JAVA)
            .setPipeline(writePipeline)
            .addParameter("runner", configuration.getRunner())
            .addParameter(
                "maxNumWorkers",
                TestProperties.getProperty("maxNumWorkers", "10", TestProperties.Type.PROPERTY))
            .build();

    return pipelineLauncher.launch(project, region, options);
  }

  private PipelineLauncher.LaunchInfo testRead() throws IOException {
    BigtableIO.Read readIO =
        BigtableIO.read()
            .withoutValidation()
            .withProjectId(project)
            .withInstanceId(resourceManager.getInstanceId())
            .withTableId(tableId);

    readPipeline
        .apply("Read from BigTable", readIO)
        .apply("Counting element", ParDo.of(new CountingFn<>(READ_ELEMENT_METRIC_NAME)));

    PipelineLauncher.LaunchConfig options =
        PipelineLauncher.LaunchConfig.builder("read-bigtable")
            .setSdk(PipelineLauncher.Sdk.JAVA)
            .setPipeline(readPipeline)
            .addParameter("runner", configuration.getRunner())
            .addParameter(
                "maxNumWorkers",
                TestProperties.getProperty("maxNumWorkers", "10", TestProperties.Type.PROPERTY))
            .build();

    return pipelineLauncher.launch(project, region, options);
  }

  /** Options for BigtableIO load test. */
  @AutoValue
  abstract static class Configuration {
    abstract Long getNumRows();

    abstract Integer getPipelineTimeout();

    abstract String getRunner();

    abstract Integer getValueSizeBytes();

    static Configuration of(long numRows, int pipelineTimeout, String runner, int valueSizeBytes) {
      return new AutoValue_BigTableIOLT_Configuration.Builder()
          .setNumRows(numRows)
          .setPipelineTimeout(pipelineTimeout)
          .setRunner(runner)
          .setValueSizeBytes(valueSizeBytes)
          .build();
    }

    @AutoValue.Builder
    abstract static class Builder {
      abstract Builder setNumRows(long numRows);

      abstract Builder setPipelineTimeout(int timeOutMinutes);

      abstract Builder setRunner(String runner);

      abstract Builder setValueSizeBytes(int valueSizeBytes);

      abstract Configuration build();
    }

    abstract Builder toBuilder();
  }

  /** Maps long number to the BigTable format record. */
  private static class MapToBigTableFormat extends DoFn<Long, KV<ByteString, Iterable<Mutation>>>
      implements Serializable {

    private final int valueSizeBytes;

    public MapToBigTableFormat(int valueSizeBytes) {
      this.valueSizeBytes = valueSizeBytes;
    }

    @ProcessElement
    public void processElement(ProcessContext c) {
      Long index = c.element();

      ByteString key = ByteString.copyFromUtf8(String.format("key%09d", index));
      Random random = new Random(index);
      byte[] valBytes = new byte[this.valueSizeBytes];
      random.nextBytes(valBytes);
      ByteString value = ByteString.copyFrom(valBytes);

      Iterable<Mutation> mutations =
          ImmutableList.of(
              Mutation.newBuilder()
                  .setSetCell(
                      Mutation.SetCell.newBuilder()
                          .setValue(value)
                          .setTimestampMicros(Instant.now().toEpochMilli() * 1000L)
                          .setFamilyName(COLUMN_FAMILY_NAME))
                  .build());
      c.output(KV.of(key, mutations));
    }
  }
}
