/*
 * -\-\-
 * 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 com.google.cloud.datastore.Batch;
import com.google.cloud.datastore.Datastore;
import com.google.cloud.datastore.DatastoreOptions;
import com.google.cloud.datastore.DatastoreReaderWriter;
import com.google.cloud.datastore.Entity;
import com.google.cloud.datastore.FullEntity;
import com.google.cloud.datastore.IncompleteKey;
import com.google.cloud.datastore.Key;
import com.google.cloud.datastore.KeyFactory;
import com.google.cloud.datastore.Query;
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 java.util.Iterator;
import java.util.List;
import java.util.Objects;

/**
 * Instrumentation for {@link Datastore} that counts operations according to https://cloud.google.com/datastore/pricing.
 */
public class InstrumentedDatastore implements Datastore, InstrumentedDatastoreReaderWriter {

  private final Stats stats;
  private final Datastore delegate;

  private InstrumentedDatastore(Datastore delegate, Stats stats) {
    this.delegate = delegate;
    this.stats = Objects.requireNonNull(stats, "stats");
  }

  @Override
  public Transaction newTransaction(TransactionOptions transactionOptions) {
    return InstrumentedTransaction.of(stats, delegate.newTransaction(transactionOptions));
  }

  @Override
  public Transaction newTransaction() {
    return InstrumentedTransaction.of(stats, delegate.newTransaction());
  }

  @Override
  public <T> T runInTransaction(TransactionCallable<T> transactionCallable) {
    return delegate.runInTransaction(rw ->
        transactionCallable.run(InstrumentedDatastoreReaderWriter.of(stats, rw)));
  }

  @Override
  public <T> T runInTransaction(TransactionCallable<T> transactionCallable, TransactionOptions transactionOptions) {
    return delegate.runInTransaction(rw ->
        transactionCallable.run(InstrumentedDatastoreReaderWriter.of(stats, rw)), transactionOptions);
  }

  @Override
  public Batch newBatch() {
    return InstrumentedBatch.of(stats, delegate.newBatch());
  }

  @Override
  public Entity get(Key key, ReadOption... readOptions) {
    stats.recordDatastoreEntityReads(key.getKind(), 1);
    return delegate.get(key, readOptions);
  }

  @Override
  public Iterator<Entity> get(Iterable<Key> keys, ReadOption... readOptions) {
    keys.forEach(key -> stats.recordDatastoreEntityReads(key.getKind(), 1));
    return delegate.get(keys, readOptions);
  }

  @Override
  public List<Entity> fetch(Iterable<Key> keys, ReadOption... readOptions) {
    keys.forEach(key -> stats.recordDatastoreEntityReads(key.getKind(), 1));
    return delegate.fetch(keys, readOptions);
  }

  @Override
  public <T> QueryResults<T> run(Query<T> query, ReadOption... readOptions) {
    final QueryResults<T> results = delegate.run(query, readOptions);
    return InstrumentedQueryResults.of(stats, query, results);
  }

  @Override
  public Key allocateId(IncompleteKey key) {
    return delegate.allocateId(key);
  }

  @Override
  public List<Key> allocateId(IncompleteKey... keys) {
    return delegate.allocateId(keys);
  }

  @Override
  public Entity add(FullEntity<?> entity) {
    return InstrumentedDatastoreReaderWriter.super.add(entity);
  }

  @Override
  public List<Entity> add(FullEntity<?>... entities) {
    return InstrumentedDatastoreReaderWriter.super.add(entities);
  }

  @Override
  public void update(Entity... entities) {
    InstrumentedDatastoreReaderWriter.super.update(entities);
  }

  @Override
  public Entity put(FullEntity<?> entity) {
    return InstrumentedDatastoreReaderWriter.super.put(entity);
  }

  @Override
  public List<Entity> put(FullEntity<?>... entities) {
    return InstrumentedDatastoreReaderWriter.super.put(entities);
  }

  @Override
  public void delete(Key... keys) {
    InstrumentedDatastoreReaderWriter.super.delete(keys);
  }

  @Override
  public KeyFactory newKeyFactory() {
    return delegate.newKeyFactory();
  }

  @Override
  public DatastoreOptions getOptions() {
    return delegate.getOptions();
  }

  @Override
  public Stats stats() {
    return stats;
  }

  @Override
  public DatastoreReaderWriter readerWriter() {
    return delegate;
  }

  public Datastore delegate() {
    return delegate;
  }

  public static InstrumentedDatastore of(Datastore delegate, Stats stats) {
    return new InstrumentedDatastore(delegate, stats);
  }
}