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

import static org.apache.beam.sdk.transforms.display.DisplayDataMatchers.hasDisplayItem;
import static org.hamcrest.Matchers.hasSize;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThat;
import static org.mockito.Matchers.any;
import static org.mockito.Mockito.argThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

import com.google.cloud.spanner.ErrorCode;
import com.google.cloud.spanner.Key;
import com.google.cloud.spanner.KeyRange;
import com.google.cloud.spanner.KeySet;
import com.google.cloud.spanner.Mutation;
import com.google.cloud.spanner.ReadOnlyTransaction;
import com.google.cloud.spanner.ResultSets;
import com.google.cloud.spanner.SpannerExceptionFactory;
import com.google.cloud.spanner.Statement;
import com.google.cloud.spanner.Struct;
import com.google.cloud.spanner.Type;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.apache.beam.sdk.testing.NeedsRunner;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import org.hamcrest.Description;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.experimental.categories.Category;
import org.junit.rules.ExpectedException;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.ArgumentMatcher;

/** Unit tests for {@link SpannerIO}. */
@RunWith(JUnit4.class)
public class SpannerIOWriteTest implements Serializable {
  private static final long CELLS_PER_KEY = 7;

  @Rule public transient TestPipeline pipeline = TestPipeline.create();
  @Rule public transient ExpectedException thrown = ExpectedException.none();

  private FakeServiceFactory serviceFactory;

  @Before
  @SuppressWarnings("unchecked")
  public void setUp() throws Exception {
    serviceFactory = new FakeServiceFactory();

    ReadOnlyTransaction tx = mock(ReadOnlyTransaction.class);
    when(serviceFactory.mockDatabaseClient().readOnlyTransaction()).thenReturn(tx);

    // Simplest schema: a table with int64 key
    preparePkMetadata(tx, Arrays.asList(pkMetadata("tEsT", "key", "ASC")));
    prepareColumnMetadata(tx, Arrays.asList(columnMetadata("tEsT", "key", "INT64", CELLS_PER_KEY)));
  }

  private static Struct columnMetadata(
      String tableName, String columnName, String type, long cellsMutated) {
    return Struct.newBuilder()
        .set("table_name")
        .to(tableName)
        .set("column_name")
        .to(columnName)
        .set("spanner_type")
        .to(type)
        .set("cells_mutated")
        .to(cellsMutated)
        .build();
  }

  private static Struct pkMetadata(String tableName, String columnName, String ordering) {
    return Struct.newBuilder()
        .set("table_name")
        .to(tableName)
        .set("column_name")
        .to(columnName)
        .set("column_ordering")
        .to(ordering)
        .build();
  }

  private void prepareColumnMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
    Type type =
        Type.struct(
            Type.StructField.of("table_name", Type.string()),
            Type.StructField.of("column_name", Type.string()),
            Type.StructField.of("spanner_type", Type.string()),
            Type.StructField.of("cells_mutated", Type.int64()));
    when(tx.executeQuery(
            argThat(
                new ArgumentMatcher<Statement>() {

                  @Override
                  public boolean matches(Object argument) {
                    if (!(argument instanceof Statement)) {
                      return false;
                    }
                    Statement st = (Statement) argument;
                    return st.getSql().contains("information_schema.columns");
                  }
                })))
        .thenReturn(ResultSets.forRows(type, rows));
  }

  private void preparePkMetadata(ReadOnlyTransaction tx, List<Struct> rows) {
    Type type =
        Type.struct(
            Type.StructField.of("table_name", Type.string()),
            Type.StructField.of("column_name", Type.string()),
            Type.StructField.of("column_ordering", Type.string()));
    when(tx.executeQuery(
            argThat(
                new ArgumentMatcher<Statement>() {

                  @Override
                  public boolean matches(Object argument) {
                    if (!(argument instanceof Statement)) {
                      return false;
                    }
                    Statement st = (Statement) argument;
                    return st.getSql().contains("information_schema.index_columns");
                  }
                })))
        .thenReturn(ResultSets.forRows(type, rows));
  }

  @Test
  public void emptyTransform() throws Exception {
    SpannerIO.Write write = SpannerIO.write();
    thrown.expect(NullPointerException.class);
    thrown.expectMessage("requires instance id to be set with");
    write.expand(null);
  }

  @Test
  public void emptyInstanceId() throws Exception {
    SpannerIO.Write write = SpannerIO.write().withDatabaseId("123");
    thrown.expect(NullPointerException.class);
    thrown.expectMessage("requires instance id to be set with");
    write.expand(null);
  }

  @Test
  public void emptyDatabaseId() throws Exception {
    SpannerIO.Write write = SpannerIO.write().withInstanceId("123");
    thrown.expect(NullPointerException.class);
    thrown.expectMessage("requires database id to be set with");
    write.expand(null);
  }

  @Test
  @Category(NeedsRunner.class)
  public void singleMutationPipeline() throws Exception {
    Mutation mutation = m(2L);
    PCollection<Mutation> mutations = pipeline.apply(Create.of(mutation));

    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory));
    pipeline.run();

    verifyBatches(batch(m(2L)));
  }

  @Test
  @Category(NeedsRunner.class)
  public void singleMutationGroupPipeline() throws Exception {
    PCollection<MutationGroup> mutations =
        pipeline.apply(Create.<MutationGroup>of(g(m(1L), m(2L), m(3L))));
    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .grouped());
    pipeline.run();

    verifyBatches(batch(m(1L), m(2L), m(3L)));
  }

  @Test
  @Category(NeedsRunner.class)
  public void batching() throws Exception {
    MutationGroup one = g(m(1L));
    MutationGroup two = g(m(2L));
    PCollection<MutationGroup> mutations = pipeline.apply(Create.of(one, two));
    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .withBatchSizeBytes(1000000000)
            .withSampler(fakeSampler(m(1000L)))
            .grouped());
    pipeline.run();

    verifyBatches(batch(m(1L), m(2L)));
  }

  @Test
  @Category(NeedsRunner.class)
  public void batchingWithDeletes() throws Exception {
    PCollection<MutationGroup> mutations =
        pipeline.apply(Create.of(g(m(1L)), g(m(2L)), g(del(3L)), g(del(4L))));
    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .withBatchSizeBytes(1000000000)
            .withSampler(fakeSampler(m(1000L)))
            .grouped());
    pipeline.run();

    verifyBatches(batch(m(1L), m(2L), del(3L), del(4L)));
  }

  @Test
  @Category(NeedsRunner.class)
  public void noBatchingRangeDelete() throws Exception {
    Mutation all = Mutation.delete("test", KeySet.all());
    Mutation prefix = Mutation.delete("test", KeySet.prefixRange(Key.of(1L)));
    Mutation range =
        Mutation.delete(
            "test", KeySet.range(KeyRange.openOpen(Key.of(1L), Key.newBuilder().build())));

    PCollection<MutationGroup> mutations =
        pipeline.apply(
            Create.of(
                g(m(1L)),
                g(m(2L)),
                g(del(5L, 6L)),
                g(delRange(50L, 55L)),
                g(delRange(11L, 20L)),
                g(all),
                g(prefix),
                g(range)));
    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .withBatchSizeBytes(1000000000)
            .withSampler(fakeSampler(m(1000L)))
            .grouped());
    pipeline.run();

    verifyBatches(
        batch(m(1L), m(2L)),
        batch(del(5L, 6L)),
        batch(delRange(11L, 20L)),
        batch(delRange(50L, 55L)),
        batch(all),
        batch(prefix),
        batch(range));
  }

  private void verifyBatches(Iterable<Mutation>... batches) {
    for (Iterable<Mutation> b : batches) {
      verify(serviceFactory.mockDatabaseClient(), times(1)).writeAtLeastOnce(mutationsInNoOrder(b));
    }
  }

  @Test
  @Category(NeedsRunner.class)
  public void sizeBatchingGroups() throws Exception {
    // Accumulate two items per batch.
    long batchSize = MutationSizeEstimator.sizeOf(g(m(1L))) * 2;

    PCollection<MutationGroup> mutations = pipeline.apply(Create.of(g(m(1L)), g(m(2L)), g(m(3L))));
    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .withBatchSizeBytes(batchSize)
            .withSampler(fakeSampler(m(1000L)))
            .grouped());

    pipeline.run();

    // The content of batches is not deterministic. Just verify that the size is correct.
    verify(serviceFactory.mockDatabaseClient(), times(1)).writeAtLeastOnce(iterableOfSize(2));
    verify(serviceFactory.mockDatabaseClient(), times(1)).writeAtLeastOnce(iterableOfSize(1));
  }

  @Test
  @Category(NeedsRunner.class)
  public void cellBatchingGroups() throws Exception {
    // Accumulate two items per batch.
    long maxNumMutations = CELLS_PER_KEY * 2;

    PCollection<MutationGroup> mutations = pipeline.apply(Create.of(g(m(1L)), g(m(2L)), g(m(3L))));
    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .withMaxNumMutations(maxNumMutations)
            .withBatchSizeBytes(Integer.MAX_VALUE)
            .withSampler(fakeSampler(m(1000L)))
            .grouped());

    pipeline.run();

    // The content of batches is not deterministic. Just verify that the size is correct.
    verify(serviceFactory.mockDatabaseClient(), times(1)).writeAtLeastOnce(iterableOfSize(2));
    verify(serviceFactory.mockDatabaseClient(), times(1)).writeAtLeastOnce(iterableOfSize(1));
  }

  @Test
  @Category(NeedsRunner.class)
  public void noBatching() throws Exception {
    PCollection<MutationGroup> mutations = pipeline.apply(Create.of(g(m(1L)), g(m(2L))));
    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .withBatchSizeBytes(1)
            .withSampler(fakeSampler(m(1000L)))
            .grouped());
    pipeline.run();

    verifyBatches(batch(m(1L)), batch(m(2L)));
  }

  @Test
  @Category(NeedsRunner.class)
  public void batchingPlusSampling() throws Exception {
    PCollection<MutationGroup> mutations =
        pipeline.apply(
            Create.of(
                g(m(1L)), g(m(2L)), g(m(3L)), g(m(4L)), g(m(5L)), g(m(6L)), g(m(7L)), g(m(8L)),
                g(m(9L)), g(m(10L))));

    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .withBatchSizeBytes(1000000000)
            .withSampler(fakeSampler(m(2L), m(5L), m(10L)))
            .grouped());
    pipeline.run();

    verifyBatches(
        batch(m(1L), m(2L)), batch(m(3L), m(4L), m(5L)), batch(m(6L), m(7L), m(8L), m(9L), m(10L)));
  }

  @Test
  @Category(NeedsRunner.class)
  public void reportFailures() throws Exception {
    PCollection<MutationGroup> mutations =
        pipeline.apply(
            Create.of(
                g(m(1L)), g(m(2L)), g(m(3L)), g(m(4L)), g(m(5L)), g(m(6L)), g(m(7L)), g(m(8L)),
                g(m(9L)), g(m(10L))));

    when(serviceFactory.mockDatabaseClient().writeAtLeastOnce(any()))
        .thenAnswer(
            invocationOnMock -> {
              throw SpannerExceptionFactory.newSpannerException(ErrorCode.ALREADY_EXISTS, "oops");
            });

    SpannerWriteResult result =
        mutations.apply(
            SpannerIO.write()
                .withProjectId("test-project")
                .withInstanceId("test-instance")
                .withDatabaseId("test-database")
                .withServiceFactory(serviceFactory)
                .withBatchSizeBytes(1000000000)
                .withFailureMode(SpannerIO.FailureMode.REPORT_FAILURES)
                .withSampler(fakeSampler(m(2L), m(5L), m(10L)))
                .grouped());
    PAssert.that(result.getFailedMutations())
        .satisfies(
            m -> {
              assertEquals(10, Iterables.size(m));
              return null;
            });
    pipeline.run();

    verifyBatches(
        batch(m(1L), m(2L)),
        batch(m(3L), m(4L), m(5L)),
        batch(m(6L), m(7L), m(8L), m(9L), m(10L)),
        // Mutations were also retried individually.
        batch(m(1L)),
        batch(m(2L)),
        batch(m(3L)),
        batch(m(4L)),
        batch(m(5L)),
        batch(m(6L)),
        batch(m(7L)),
        batch(m(8L)),
        batch(m(9L)),
        batch(m(10L)));
  }

  @Test
  @Category(NeedsRunner.class)
  public void noBatchingPlusSampling() throws Exception {
    PCollection<MutationGroup> mutations =
        pipeline.apply(Create.of(g(m(1L)), g(m(2L)), g(m(3L)), g(m(4L)), g(m(5L))));
    mutations.apply(
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withServiceFactory(serviceFactory)
            .withBatchSizeBytes(1)
            .withSampler(fakeSampler(m(2L)))
            .grouped());

    pipeline.run();

    verifyBatches(batch(m(1L)), batch(m(2L)), batch(m(3L)), batch(m(4L)), batch(m(5L)));
  }

  @Test
  public void displayData() throws Exception {
    SpannerIO.Write write =
        SpannerIO.write()
            .withProjectId("test-project")
            .withInstanceId("test-instance")
            .withDatabaseId("test-database")
            .withBatchSizeBytes(123);

    DisplayData data = DisplayData.from(write);
    assertThat(data.items(), hasSize(4));
    assertThat(data, hasDisplayItem("projectId", "test-project"));
    assertThat(data, hasDisplayItem("instanceId", "test-instance"));
    assertThat(data, hasDisplayItem("databaseId", "test-database"));
    assertThat(data, hasDisplayItem("batchSizeBytes", 123));
  }

  private static MutationGroup g(Mutation m, Mutation... other) {
    return MutationGroup.create(m, other);
  }

  private static Mutation m(Long key) {
    return Mutation.newInsertOrUpdateBuilder("test").set("key").to(key).build();
  }

  private static Iterable<Mutation> batch(Mutation... m) {
    return Arrays.asList(m);
  }

  private static Mutation del(Long... keys) {

    KeySet.Builder builder = KeySet.newBuilder();
    for (Long key : keys) {
      builder.addKey(Key.of(key));
    }
    return Mutation.delete("test", builder.build());
  }

  private static Mutation delRange(Long start, Long end) {
    return Mutation.delete("test", KeySet.range(KeyRange.closedClosed(Key.of(start), Key.of(end))));
  }

  private static Iterable<Mutation> mutationsInNoOrder(Iterable<Mutation> expected) {
    final ImmutableSet<Mutation> mutations = ImmutableSet.copyOf(expected);
    return argThat(
        new ArgumentMatcher<Iterable<Mutation>>() {

          @Override
          public boolean matches(Object argument) {
            if (!(argument instanceof Iterable)) {
              return false;
            }
            ImmutableSet<Mutation> actual = ImmutableSet.copyOf((Iterable) argument);
            return actual.equals(mutations);
          }

          @Override
          public void describeTo(Description description) {
            description.appendText("Iterable must match ").appendValue(mutations);
          }
        });
  }

  private Iterable<Mutation> iterableOfSize(final int size) {
    return argThat(
        new ArgumentMatcher<Iterable<Mutation>>() {

          @Override
          public boolean matches(Object argument) {
            return argument instanceof Iterable && Iterables.size((Iterable<?>) argument) == size;
          }

          @Override
          public void describeTo(Description description) {
            description.appendText("The size of the iterable must equal ").appendValue(size);
          }
        });
  }

  private static FakeSampler fakeSampler(Mutation... mutations) {
    SpannerSchema.Builder schema = SpannerSchema.builder();
    schema.addColumn("test", "key", "INT64", CELLS_PER_KEY);
    schema.addKeyPart("test", "key", false);
    return new FakeSampler(schema.build(), Arrays.asList(mutations));
  }

  private static class FakeSampler
      extends PTransform<PCollection<KV<String, byte[]>>, PCollection<KV<String, List<byte[]>>>> {

    private final SpannerSchema schema;
    private final List<Mutation> mutations;

    private FakeSampler(SpannerSchema schema, List<Mutation> mutations) {
      this.schema = schema;
      this.mutations = mutations;
    }

    @Override
    public PCollection<KV<String, List<byte[]>>> expand(PCollection<KV<String, byte[]>> input) {
      MutationGroupEncoder coder = new MutationGroupEncoder(schema);
      Map<String, List<byte[]>> map = new HashMap<>();
      for (Mutation m : mutations) {
        String table = m.getTable().toLowerCase();
        List<byte[]> list = map.computeIfAbsent(table, k -> new ArrayList<>());
        list.add(coder.encodeKey(m));
      }
      List<KV<String, List<byte[]>>> result = new ArrayList<>();
      for (Map.Entry<String, List<byte[]>> entry : map.entrySet()) {
        entry.getValue().sort(SpannerIO.SerializableBytesComparator.INSTANCE);
        result.add(KV.of(entry.getKey(), entry.getValue()));
      }
      return input.getPipeline().apply(Create.of(result));
    }
  }
}
