/*
 * -\-\-
 * Spotify Styx Scheduler Service
 * --
 * Copyright (C) 2018 Spotify AB
 * --
 * Licensed 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 com.spotify.styx.storage;

import static com.google.datastore.v1.QueryResultBatch.MoreResultsType.MORE_RESULTS_AFTER_CURSOR;
import static java.util.Arrays.asList;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.instanceOf;
import static org.hamcrest.Matchers.is;
import static org.hamcrest.Matchers.sameInstance;
import static org.junit.Assert.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;

import com.google.cloud.datastore.BaseEntity;
import com.google.cloud.datastore.Batch;
import com.google.cloud.datastore.Datastore;
import com.google.cloud.datastore.Datastore.TransactionCallable;
import com.google.cloud.datastore.DatastoreBatchWriter;
import com.google.cloud.datastore.DatastoreOptions;
import com.google.cloud.datastore.DatastoreReader;
import com.google.cloud.datastore.DatastoreReaderWriter;
import com.google.cloud.datastore.DatastoreWriter;
import com.google.cloud.datastore.Entity;
import com.google.cloud.datastore.EntityQuery;
import com.google.cloud.datastore.IncompleteKey;
import com.google.cloud.datastore.Key;
import com.google.cloud.datastore.KeyFactory;
import com.google.cloud.datastore.KeyQuery;
import com.google.cloud.datastore.ProjectionEntity;
import com.google.cloud.datastore.ProjectionEntityQuery;
import com.google.cloud.datastore.QueryResults;
import com.google.cloud.datastore.ReadOption;
import com.google.cloud.datastore.Transaction;
import com.google.datastore.v1.TransactionOptions;
import com.spotify.styx.monitoring.Stats;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mock;
import org.mockito.junit.MockitoJUnitRunner;

@SuppressWarnings({"ResultOfMethodCallIgnored", "unchecked"})
@RunWith(MockitoJUnitRunner.class)
public class InstrumentedDatastoreTest {

  private static final String TEST_PROJECT = "test-project";
  private static final String TEST_KIND = "test-kind";
  private static final String TEST_KIND_1 = "test-kind-1";
  private static final String TEST_KIND_2 = "test-kind-2";
  private static final Key TEST_KEY = Key.newBuilder(TEST_PROJECT, TEST_KIND, "test").build();
  private static final Key TEST_KEY_1 = Key.newBuilder(TEST_PROJECT, TEST_KIND_1, "test-1").build();
  private static final Key TEST_KEY_2 = Key.newBuilder(TEST_PROJECT, TEST_KIND_2, "test-2").build();
  private static final Entity TEST_ENTITY = Entity.newBuilder(TEST_KEY).build();
  private static final Entity TEST_ENTITY_1 = Entity.newBuilder(TEST_KEY_1).build();
  private static final Entity TEST_ENTITY_2 = Entity.newBuilder(TEST_KEY_2).build();
  private static final EntityQuery TEST_ENTITY_QUERY = EntityQuery.newEntityQueryBuilder().setKind(TEST_KIND).build();
  private static final KeyQuery TEST_KEY_QUERY = EntityQuery.newKeyQueryBuilder().setKind(TEST_KIND).build();
  private static final ProjectionEntityQuery TEST_PROJECTION_QUERY_WITHOUT_DISTINCT =
      EntityQuery.newProjectionEntityQueryBuilder().setKind(TEST_KIND).build();
  private static final ProjectionEntityQuery TEST_PROJECTION_QUERY_WITH_DISTINCT =
      EntityQuery.newProjectionEntityQueryBuilder().setKind(TEST_KIND).addDistinctOn("foo", "bar").build();
  private static final KeyFactory KEY_FACTORY = new KeyFactory(TEST_PROJECT);

  private InstrumentedDatastore instrumentedDatastore;

  @Mock Datastore datastore;
  @Mock Stats stats;
  @Mock QueryResults<Entity> entityQueryResults;
  @Mock QueryResults<Key> keyQueryResults;
  @Mock QueryResults<ProjectionEntity> projectionQueryResults;
  @Mock Transaction transaction;
  @Mock Batch batch;
  @Mock ReadOption readOption1;
  @Mock ReadOption readOption2;
  @Mock IncompleteKey incompleteKey1;
  @Mock IncompleteKey incompleteKey2;
  @Mock DatastoreOptions options;
  @Mock BaseEntity<?> entity;

  @Before
  public void setUp() throws Exception {
    instrumentedDatastore = InstrumentedDatastore.of(datastore, stats);
  }

  @Test
  public void testDatastore() {
    testReaderWriter(instrumentedDatastore, datastore);

    // Entity get(Key key, ReadOption... options);
    instrumentedDatastore.get(TEST_KEY, readOption1, readOption2);
    verify(datastore).get(TEST_KEY, readOption1, readOption2);
    verify(stats).recordDatastoreEntityReads(TEST_KIND, 1);
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore);

    // Iterator<Entity> get(Iterable<Key> keys, ReadOption... options);
    instrumentedDatastore.get(asList(TEST_KEY_1, TEST_KEY_2), readOption1, readOption2);
    verify(datastore).get(asList(TEST_KEY_1, TEST_KEY_2), readOption1, readOption2);
    verify(stats).recordDatastoreEntityReads(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityReads(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore);

    // List<Entity> fetch(Iterable<Key> keys, ReadOption... options);
    instrumentedDatastore.fetch(asList(TEST_KEY_1, TEST_KEY_2), readOption1, readOption2);
    verify(datastore).fetch(asList(TEST_KEY_1, TEST_KEY_2), readOption1, readOption2);
    verify(stats).recordDatastoreEntityReads(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityReads(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore);

    // <T> QueryResults<T> run(Query<T> query, ReadOption... options);

    // Entity Query
    when(datastore.run(TEST_ENTITY_QUERY, readOption1, readOption2)).thenReturn(entityQueryResults);
    var instrumentedEntityQueryResults =
        instrumentedDatastore.run(TEST_ENTITY_QUERY, readOption1, readOption2);
    verify(datastore).run(TEST_ENTITY_QUERY, readOption1, readOption2);
    verify(stats).recordDatastoreQueries(TEST_KIND, 1);
    testInstrumentedQueryResults(instrumentedEntityQueryResults, entityQueryResults);
    verifyNoMoreInteractions(stats, datastore, entityQueryResults);
    reset(stats, datastore, entityQueryResults);

    // Key Query
    when(datastore.run(TEST_KEY_QUERY, readOption1, readOption2)).thenReturn(keyQueryResults);
    assertThat(instrumentedDatastore.run(TEST_KEY_QUERY, readOption1, readOption2), sameInstance(keyQueryResults));
    verify(datastore).run(TEST_KEY_QUERY, readOption1, readOption2);
    verify(stats).recordDatastoreQueries(TEST_KIND, 1);
    verifyNoMoreInteractions(stats, datastore, keyQueryResults);
    reset(stats, datastore, keyQueryResults);

    // Projection Query - With Distinct
    when(datastore.run(TEST_PROJECTION_QUERY_WITH_DISTINCT, readOption1, readOption2))
        .thenReturn(projectionQueryResults);
    var instrumentedProjectionQueryResults =
        instrumentedDatastore.run(TEST_PROJECTION_QUERY_WITH_DISTINCT, readOption1, readOption2);
    verify(datastore).run(TEST_PROJECTION_QUERY_WITH_DISTINCT, readOption1, readOption2);
    verify(stats).recordDatastoreQueries(TEST_KIND, 1);
    testInstrumentedQueryResults(instrumentedProjectionQueryResults, projectionQueryResults);
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore, projectionQueryResults);

    // Project Query - Without Distinct
    when(datastore.run(TEST_PROJECTION_QUERY_WITHOUT_DISTINCT, readOption1, readOption2))
        .thenReturn(projectionQueryResults);
    assertThat(instrumentedDatastore.run(TEST_PROJECTION_QUERY_WITHOUT_DISTINCT, readOption1, readOption2),
        is(sameInstance(projectionQueryResults)));
    verify(stats).recordDatastoreQueries(TEST_KIND, 1);
    verify(datastore).run(TEST_PROJECTION_QUERY_WITHOUT_DISTINCT, readOption1, readOption2);
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore);

    // Key allocateId(IncompleteKey key);
    when(instrumentedDatastore.allocateId(incompleteKey1)).thenReturn(TEST_KEY_1);
    assertThat(instrumentedDatastore.allocateId(incompleteKey1), is(TEST_KEY_1));
    verify(datastore).allocateId(incompleteKey1);
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore);

    // List<Key> allocateId(IncompleteKey... keys);
    when(instrumentedDatastore.allocateId(incompleteKey1, incompleteKey2)).thenReturn(asList(TEST_KEY_1, TEST_KEY_2));
    assertThat(instrumentedDatastore.allocateId(incompleteKey1, incompleteKey2), contains(TEST_KEY_1, TEST_KEY_2));
    verify(datastore).allocateId(incompleteKey1, incompleteKey2);
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore);

    // KeyFactory newKeyFactory();
    when(datastore.newKeyFactory()).thenReturn(KEY_FACTORY);
    assertThat(instrumentedDatastore.newKeyFactory(), is(KEY_FACTORY));
    verify(datastore).newKeyFactory();
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore);

    // OptionsT getOptions();
    when(datastore.getOptions()).thenReturn(options);
    assertThat(instrumentedDatastore.getOptions(), is(options));
    verify(datastore).getOptions();
    verifyNoMoreInteractions(stats, datastore);
    reset(stats, datastore);
  }

  @Test
  public void newTransaction() {
    // Transaction newTransaction(TransactionOptions options);
    when(datastore.newTransaction(any(TransactionOptions.class))).thenReturn(transaction);
    var transactionOptions = TransactionOptions.newBuilder().build();
    var instrumented1 = instrumentedDatastore.newTransaction(transactionOptions);
    verify(datastore).newTransaction(transactionOptions);
    verifyNoMoreInteractions(datastore);
    reset(datastore);
    testTransaction(instrumented1);

    // Transaction newTransaction();
    when(datastore.newTransaction()).thenReturn(transaction);
    var instrumented2 = instrumentedDatastore.newTransaction();
    verify(datastore).newTransaction();
    verifyNoMoreInteractions(datastore);
    reset(datastore);
    testTransaction(instrumented2);
  }

  @Test
  public void runInTransaction() {
    when(datastore.runInTransaction(any())).then(a ->
        a.<TransactionCallable<String>>getArgument(0)
            .run(transaction));

    var foobar = instrumentedDatastore.runInTransaction(tx -> {
      testReaderWriter(tx, transaction);
      return "foobar";
    });

    verify(datastore).runInTransaction(any());
    verifyNoMoreInteractions(datastore);

    assertThat(foobar, is("foobar"));
  }

  @Test
  public void runInTransactionWithOptions() {
    when(datastore.runInTransaction(any(), any())).then(a ->
        a.<TransactionCallable<String>>getArgument(0)
            .run(transaction));

    var transactionOptions = TransactionOptions.newBuilder().build();

    var foobar = instrumentedDatastore.runInTransaction(tx -> {
      testReaderWriter(tx, transaction);
      return "foobar";
    }, transactionOptions);

    verify(datastore).runInTransaction(any(), eq(transactionOptions));
    verifyNoMoreInteractions(datastore);

    assertThat(foobar, is("foobar"));
  }

  @Test
  public void newBatch() {
    when(datastore.newBatch()).thenReturn(batch);
    var instrumentedBatch = instrumentedDatastore.newBatch();
    verify(datastore).newBatch();
    verifyNoMoreInteractions(datastore);
    reset(datastore);
    testBatchWriter(instrumentedBatch, batch);

    instrumentedBatch.submit();
    verify(batch).submit();
    verifyNoMoreInteractions(batch);
    reset(batch);

    when(batch.getDatastore()).thenReturn(datastore);
    var instrumentedDatastore = instrumentedBatch.getDatastore();
    assertThat(instrumentedDatastore, instanceOf(InstrumentedDatastore.class));
    assertThat(((InstrumentedDatastore) instrumentedDatastore).delegate(), is(datastore));
    verify(batch).getDatastore();
    verifyNoMoreInteractions(batch);
  }

  private void testTransaction(Transaction instrumented) {
    testReaderWriter(instrumented, transaction);

    // void putWithDeferredIdAllocation(FullEntity<?>... entities);
    instrumented.putWithDeferredIdAllocation(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(transaction).putWithDeferredIdAllocation(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats);
    reset(stats);

    // boolean isActive();
    when(transaction.isActive()).thenReturn(true);
    assertThat(instrumented.isActive(), is(true));
    verify(transaction).isActive();

    // void rollback();
    instrumented.rollback();
    verify(transaction).rollback();

    // Response commit();
    instrumented.commit();
    verify(transaction).commit();

    // Datastore getDatastore();
    when(transaction.getDatastore()).thenReturn(datastore);
    var instrumentedDatastore = instrumented.getDatastore();
    assertThat(instrumentedDatastore, instanceOf(InstrumentedDatastore.class));
    assertThat(((InstrumentedDatastore) instrumentedDatastore).delegate(), is(datastore));
    verify(transaction).getDatastore();

    verifyNoMoreInteractions(transaction);
    reset(transaction);
  }

  private void testReaderWriter(DatastoreReaderWriter instrumented, DatastoreReaderWriter delegate) {
    testReader(instrumented, delegate);
    testWriter(instrumented, delegate);
  }

  private void testWriter(DatastoreWriter instrumented, DatastoreWriter delegate) {

    // Entity add(FullEntity<?> entity);
    instrumented.add(TEST_ENTITY);
    verify(delegate).add(TEST_ENTITY);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // List<Entity> add(FullEntity<?>... entities);
    instrumented.add(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(delegate).add(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // void update(Entity... entities);
    instrumented.update(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(delegate).update(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // Entity put(FullEntity<?> entity);
    instrumented.put(TEST_ENTITY);
    verify(delegate).put(TEST_ENTITY);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // List<Entity> put(FullEntity<?>... entities);
    instrumented.put(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(delegate).put(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // void delete(Key... keys);
    instrumented.delete(TEST_KEY_1, TEST_KEY_2);
    verify(delegate).delete(TEST_KEY_1, TEST_KEY_2);
    verify(stats).recordDatastoreEntityDeletes(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityDeletes(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);
  }

  private void testReader(DatastoreReader instrumented, DatastoreReader delegate) {
    // Entity get(Key key);
    instrumented.get(TEST_KEY);
    verify(delegate).get(TEST_KEY);
    verify(stats).recordDatastoreEntityReads(TEST_KIND, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // Iterator<Entity> get(Key... keys);
    instrumented.get(TEST_KEY_1, TEST_KEY_2);
    verify(delegate).get(TEST_KEY_1, TEST_KEY_2);
    verify(stats).recordDatastoreEntityReads(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityReads(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // List<Entity> fetch(Key... keys);
    instrumented.fetch(TEST_KEY_1, TEST_KEY_2);
    verify(delegate).fetch(TEST_KEY_1, TEST_KEY_2);
    verify(stats).recordDatastoreEntityReads(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityReads(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // <T> QueryResults<T> run(Query<T> query);

    // Entity Query
    when(delegate.run(TEST_ENTITY_QUERY)).thenReturn(entityQueryResults);
    var instrumentedEntityQueryResults = instrumented.run(TEST_ENTITY_QUERY);
    verify(delegate).run(TEST_ENTITY_QUERY);
    verify(stats).recordDatastoreQueries(TEST_KIND, 1);
    testInstrumentedQueryResults(instrumentedEntityQueryResults, entityQueryResults);
    verifyNoMoreInteractions(stats, delegate, entityQueryResults);
    reset(stats, delegate, entityQueryResults);

    // Key Query
    when(delegate.run(TEST_KEY_QUERY)).thenReturn(keyQueryResults);
    assertThat(instrumented.run(TEST_KEY_QUERY), sameInstance(keyQueryResults));
    verify(delegate).run(TEST_KEY_QUERY);
    verify(stats).recordDatastoreQueries(TEST_KIND, 1);
    verifyNoMoreInteractions(stats, delegate, keyQueryResults);
    reset(stats, delegate, keyQueryResults);

    // Projection Query - With Distinct
    when(delegate.run(TEST_PROJECTION_QUERY_WITH_DISTINCT)).thenReturn(projectionQueryResults);
    var instrumentedProjectionQueryResults =
        instrumented.run(TEST_PROJECTION_QUERY_WITH_DISTINCT);
    verify(delegate).run(TEST_PROJECTION_QUERY_WITH_DISTINCT);
    verify(stats).recordDatastoreQueries(TEST_KIND, 1);
    testInstrumentedQueryResults(instrumentedProjectionQueryResults, projectionQueryResults);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate, projectionQueryResults);

    // Project Query - Without Distinct
    when(delegate.run(TEST_PROJECTION_QUERY_WITHOUT_DISTINCT)).thenReturn(projectionQueryResults);
    assertThat(instrumented.run(TEST_PROJECTION_QUERY_WITHOUT_DISTINCT), is(sameInstance(projectionQueryResults)));
    verify(stats).recordDatastoreQueries(TEST_KIND, 1);
    verify(delegate).run(TEST_PROJECTION_QUERY_WITHOUT_DISTINCT);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);
  }

  private <T extends BaseEntity<?>> void testInstrumentedQueryResults(QueryResults<T> instrumented,
      QueryResults<T> delegate) {
    when(delegate.hasNext()).thenReturn(true);
    when(delegate.next()).thenReturn((T) entity);
    when(delegate.getSkippedResults()).thenReturn(17);
    when(delegate.getMoreResults()).thenReturn(MORE_RESULTS_AFTER_CURSOR);
    assertThat(instrumented.hasNext(), is(true));
    verify(delegate).hasNext();
    assertThat(instrumented.next(), is(entity));
    verify(delegate).next();
    assertThat(instrumented.getSkippedResults(), is(17));
    assertThat(instrumented.getMoreResults(), is(MORE_RESULTS_AFTER_CURSOR));
    verify(delegate).getSkippedResults();
    verify(delegate).getMoreResults();
    verify(stats).recordDatastoreEntityReads(TEST_KIND, 1);
  }

  private void testBatchWriter(DatastoreBatchWriter instrumented, DatastoreBatchWriter delegate) {
    testWriter(instrumented, delegate);

    // void addWithDeferredIdAllocation(FullEntity<?>... entities);
    instrumented.addWithDeferredIdAllocation(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(delegate).addWithDeferredIdAllocation(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // void putWithDeferredIdAllocation(FullEntity<?>... entities);
    instrumented.putWithDeferredIdAllocation(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(delegate).putWithDeferredIdAllocation(TEST_ENTITY_1, TEST_ENTITY_2);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_1, 1);
    verify(stats).recordDatastoreEntityWrites(TEST_KIND_2, 1);
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);

    // boolean isActive();
    when(delegate.isActive()).thenReturn(true);
    assertThat(instrumented.isActive(), is(true));
    verify(delegate).isActive();
    verifyNoMoreInteractions(stats, delegate);
    reset(stats, delegate);
  }
}
